#!/usr/bin/env python3
"""End-to-end pipeline: universe -> signal -> cross-sectional IC
-> multi-stock backtest (AlphaStrategy + RankEqualWeightBuilder) -> reports.

Usage:
    python3 run_example.py --universe hs300 --signal reversal
    python3 run_example.py --universe csi500 --signal reversal_vol
"""
import argparse
import logging

import pandas as pd

from analysis.report import generate_report
from backtest.config import BacktestConfig
from backtest.runner import BacktestRunner
from data.downloader import download_batch
from data.universe import SYMBOLS, CSI500_SYMBOLS
from eval.metrics import evaluate_cross_sectional
from portfolio.builder import RankEqualWeightBuilder
from signals.reversal import ReversalSignal
from signals.reversal_vol import ReversalVolSignal
from strategies.alpha_strategy import AlphaStrategy

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


def _forward_returns(data: dict[str, pd.DataFrame], horizon: int) -> pd.DataFrame:
    """Build a date-indexed DataFrame of ``horizon``-day forward returns per stock."""
    forward_returns: dict[str, pd.Series] = {}
    for sym, df in data.items():
        fwd = df["close"].pct_change(horizon).shift(-horizon)
        fwd.index = pd.to_datetime(df["date"])
        forward_returns[sym] = fwd
    return pd.DataFrame(forward_returns)


def main(forward_horizon: int = 5, universe: str = "csi500", signal_name: str = "reversal_vol"):
    universes = {"hs300": SYMBOLS, "csi500": CSI500_SYMBOLS}
    symbols = universes.get(universe, CSI500_SYMBOLS)[:30]

    signals = {
        "reversal": ReversalSignal(lookback=5),
        "reversal_vol": ReversalVolSignal(lookback=5, vol_window=20),
    }
    signal = signals.get(signal_name, ReversalVolSignal(lookback=5, vol_window=20))

    start, end = "2023-01-01", "2024-12-31"
    initial_cash = 1_000_000

    logger.info(f"Universe: {universe} ({len(symbols)} stocks), Signal: {signal.name}")

    # 1-2. Download daily data for the universe.
    data = download_batch(symbols, start, end)
    data = {s: df for s, df in data.items() if df is not None and not df.empty}
    logger.info(f"Downloaded {len(data)}/{len(symbols)} symbols")

    # 3. Compute the signal per stock.
    signal_series: dict[str, pd.Series] = {}
    for sym, df in data.items():
        sig = signal.compute(df)
        sig.index = pd.to_datetime(df["date"])
        signal_series[sym] = sig

    # 4. Cross-sectional IC at the matching forward horizon.
    signals_df = pd.DataFrame(signal_series)
    returns_df = _forward_returns(data, forward_horizon)
    signal_eval = evaluate_cross_sectional(signals_df, returns_df)

    # 4b. Multi-horizon IC.
    horizon_evals = {
        h: evaluate_cross_sectional(signals_df, _forward_returns(data, h))
        for h in (1, 5, 20)
    }

    # 5. Attach the signal column to each DataFrame and build feeds.
    config = BacktestConfig(
        symbols=list(data.keys()),
        start_date=start,
        end_date=end,
        initial_cash=initial_cash,
        sizer_percent=0.95,
    )
    runner = BacktestRunner(config)
    builder = RankEqualWeightBuilder(top_n=5)
    for sym, df in data.items():
        df = df.copy()
        df["signal"] = signal.compute(df).values
        runner.add_signal_data(df, name=sym)

    # 6. Run the multi-stock backtest.
    results = runner.run_prepared(AlphaStrategy, {"builder": builder})

    # 7. Reports.
    artifacts = generate_report(
        results, signal_eval, output_dir="reports/", initial_cash=initial_cash
    )

    # 8. Print summary.
    print("\nSIGNAL IC")
    print("=" * 50)
    print(f"Universe: {universe}  |  Signal: {signal.name}")
    print(f"IC mean / std / IR:       {signal_eval['ic_mean']:.4f} / "
          f"{signal_eval['ic_std']:.4f} / {signal_eval['ir']:.4f}")
    print(f"Rank IC mean / std / IR:  {signal_eval['rank_ic_mean']:.4f} / "
          f"{signal_eval['rank_ic_std']:.4f} / {signal_eval['rank_ir']:.4f}")
    print(f"Hit rate:                 {signal_eval['hit_rate']:.2%}")
    print(f"Periods:                  {signal_eval['n_periods']}")

    print("\nMULTI-HORIZON IC")
    print("=" * 50)
    print(f"{'Horizon':>8}  {'Rank IC':>9}  {'Rank IR':>9}  {'Hit rate':>9}  {'Periods':>8}")
    for h, ev in horizon_evals.items():
        print(f"{f'{h}d':>8}  {ev['rank_ic_mean']:>9.4f}  {ev['rank_ir']:>9.4f}  "
              f"{ev['hit_rate']:>8.2%}  {ev['n_periods']:>8}")

    print(f"\nReports written to: {artifacts}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chinese equity quant backtest")
    parser.add_argument("--universe", default="csi500", choices=["hs300", "csi500"])
    parser.add_argument("--signal", default="reversal_vol", choices=["reversal", "reversal_vol"])
    args = parser.parse_args()
    main(universe=args.universe, signal_name=args.signal)