"""Does V3 regime timing + S&P 500 stock picking improve over either alone?

Variants tested:
  1. RegimeStockPicker top-10  — V3 regime, risk-on = top-10 momentum stocks
  2. RegimeStockPicker top-20  — V3 regime, risk-on = top-20 momentum stocks
  3. RegimeRecovery   top-10   — V3 regime, risk-on = recovery+momentum top-10
  4. RecoveryMomentum top-10   — pure stock picker, no regime filter (baseline)
  5. TrendRider V7             — leveraged ETFs (current SOTA)
  6. SPY buy-and-hold          — benchmark
"""
from __future__ import annotations

import sys

sys.path.insert(0, ".")

import numpy as np
import pandas as pd

import data_manager
import metrics
import universe_history as uh
from main import backtest
from strategies.base import Strategy
from strategies.permanent import TrendRiderV3
from strategies.recovery_momentum import RecoveryMomentumStrategy
from strategies.trend_rider_v7 import TrendRiderV7
from universe import UNIVERSES

YEARS = 10
CAPITAL = 100_000
TX_COST = 0.001
FIXED_FEE = 2.0


# ---------------------------------------------------------------------------
# Strategy: V3 regime gate + cross-sectional momentum on S&P 500
# ---------------------------------------------------------------------------

class RegimeStockPicker(Strategy):
    """V3 macro regime + S&P 500 momentum stock selection.

    Risk-on: equal-weight top-N stocks by ``mom_lookback``-day momentum.
    Risk-off: momentum leader of (GLD, DBC).
    """

    def __init__(
        self,
        stock_tickers: list[str],
        top_n: int = 10,
        signal: str = "SPY",
        defensive: tuple[str, ...] = ("GLD", "DBC"),
        ma_long: int = 150,
        mom_lookback: int = 63,
        rebal_every: int = 21,
    ):
        self.stock_tickers = stock_tickers
        self.top_n = top_n
        self.signal = signal
        self.defensive = defensive
        self.ma_long = ma_long
        self.mom_lookback = mom_lookback
        self.rebal_every = rebal_every
        self._v3 = TrendRiderV3(
            signal=signal, risk_on=("TQQQ", "UPRO"), risk_off=defensive,
            ma_long=ma_long,
        )

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        w = pd.DataFrame(np.nan, index=data.index, columns=data.columns)

        if self.signal not in data.columns:
            return w.fillna(0.0)

        sig_arr = data[self.signal].to_numpy()
        mom = data.pct_change(self.mom_lookback, fill_method=None)

        avail_stocks = [t for t in self.stock_tickers if t in data.columns]
        avail_def = [t for t in self.defensive if t in data.columns]

        need = max(self.ma_long, self.mom_lookback + 1,
                   self._v3.vol_window + 1, self._v3.dd_window,
                   self._v3.peak_window) + 1

        regime: str | None = None
        bars = 0

        for i in range(len(data)):
            if i < need:
                continue

            closes = sig_arr[:i]
            if np.isnan(closes[-1]):
                continue

            desired = self._v3._desired_regime(closes, regime)
            changed = False
            if regime is None:
                regime, bars, changed = desired, 0, True
            else:
                bars += 1
                if desired != regime and bars >= 15:
                    regime, bars, changed = desired, 0, True

            if not changed and (i - need) % self.rebal_every != 0:
                continue

            row = {c: 0.0 for c in data.columns}
            dt = data.index[i]

            if regime == "risk_on":
                m = mom.iloc[i][avail_stocks].dropna()
                valid = m.index[data.loc[dt, m.index].notna()]
                m = m[valid]
                m = m[m > 0]
                top = m.nlargest(min(self.top_n, len(m)))
                if len(top) > 0:
                    wt = 1.0 / len(top)
                    for t in top.index:
                        row[t] = wt
                elif avail_def:
                    row[avail_def[0]] = 1.0
            else:
                if avail_def:
                    dm = mom.iloc[i][avail_def].dropna()
                    best = dm.idxmax() if len(dm) > 0 else avail_def[0]
                    row[best] = 1.0

            for c, v in row.items():
                w.at[dt, c] = v

        w = w.ffill().fillna(0.0)
        return w.shift(1).fillna(0.0)


# ---------------------------------------------------------------------------
# Strategy: V3 regime gate + recovery-momentum composite on S&P 500
# ---------------------------------------------------------------------------

class RegimeRecoveryPicker(Strategy):
    """V3 regime + recovery-momentum composite for stock selection.

    Uses the same factor as RecoveryMomentumStrategy but only during risk-on.
    """

    def __init__(
        self,
        stock_tickers: list[str],
        top_n: int = 10,
        signal: str = "SPY",
        defensive: tuple[str, ...] = ("GLD", "DBC"),
        ma_long: int = 150,
        recovery_window: int = 63,
        mom_lookback: int = 252,
        mom_skip: int = 21,
        rebal_every: int = 21,
    ):
        self.stock_tickers = stock_tickers
        self.top_n = top_n
        self.signal = signal
        self.defensive = defensive
        self.ma_long = ma_long
        self.recovery_window = recovery_window
        self.mom_lookback = mom_lookback
        self.mom_skip = mom_skip
        self.rebal_every = rebal_every
        self._v3 = TrendRiderV3(
            signal=signal, risk_on=("TQQQ", "UPRO"), risk_off=defensive,
            ma_long=ma_long,
        )

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        w = pd.DataFrame(np.nan, index=data.index, columns=data.columns)

        if self.signal not in data.columns:
            return w.fillna(0.0)

        sig_arr = data[self.signal].to_numpy()

        stock_data = data[[t for t in self.stock_tickers if t in data.columns]]
        recovery = stock_data / stock_data.rolling(self.recovery_window).min() - 1
        momentum = stock_data.shift(self.mom_skip).pct_change(
            self.mom_lookback - self.mom_skip, fill_method=None,
        )
        rec_rank = recovery.rank(axis=1, pct=True, na_option="keep")
        mom_rank = momentum.rank(axis=1, pct=True, na_option="keep")
        composite = 0.5 * rec_rank + 0.5 * mom_rank
        stock_rank = composite.rank(axis=1, ascending=False, na_option="bottom")

        def_mom = data[[t for t in self.defensive if t in data.columns]].pct_change(63, fill_method=None)
        avail_def = [t for t in self.defensive if t in data.columns]

        need = max(self.ma_long, self.mom_lookback + 1,
                   self._v3.vol_window + 1, self._v3.dd_window,
                   self._v3.peak_window, self.recovery_window) + 1

        regime: str | None = None
        bars = 0

        for i in range(len(data)):
            if i < need:
                continue
            closes = sig_arr[:i]
            if np.isnan(closes[-1]):
                continue

            desired = self._v3._desired_regime(closes, regime)
            changed = False
            if regime is None:
                regime, bars, changed = desired, 0, True
            else:
                bars += 1
                if desired != regime and bars >= 15:
                    regime, bars, changed = desired, 0, True

            if not changed and (i - need) % self.rebal_every != 0:
                continue

            row = {c: 0.0 for c in data.columns}
            dt = data.index[i]

            if regime == "risk_on":
                ranks_i = stock_rank.iloc[i]
                n_valid = composite.iloc[i].notna().sum()
                if n_valid >= self.top_n:
                    top = ranks_i[ranks_i <= self.top_n].index
                    if len(top) > 0:
                        wt = 1.0 / len(top)
                        for t in top:
                            row[t] = wt
                if sum(row.values()) < 0.01 and avail_def:
                    row[avail_def[0]] = 1.0
            else:
                if avail_def:
                    dm = def_mom.iloc[i][avail_def].dropna()
                    best = dm.idxmax() if len(dm) > 0 else avail_def[0]
                    row[best] = 1.0

            for c, v in row.items():
                w.at[dt, c] = v

        w = w.ffill().fillna(0.0)
        return w.shift(1).fillna(0.0)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    print("=" * 80)
    print("  REGIME + STOCK PICKER EVALUATION")
    print("=" * 80)

    # --- Load S&P 500 data (PIT-safe) ---
    print("\n[1/3] Loading S&P 500 universe (PIT)...")
    universe = UNIVERSES["us"]
    tickers = universe["fetch"]()
    pit_intervals = uh.load_sp500_history()
    hist_tickers = uh.all_tickers_ever(pit_intervals)

    etf_extra = ["SPY", "GLD", "DBC", "SHY", "TQQQ", "UPRO", "TLT"]
    all_tickers = sorted(set(tickers + hist_tickers + etf_extra))
    print(f"   Downloading {len(all_tickers)} tickers...")

    data = data_manager.update("us", all_tickers, with_open=False)
    if isinstance(data, tuple):
        data = data[0]

    cutoff = data.index[-1] - pd.DateOffset(years=YEARS)
    data = data[data.index >= cutoff]
    data = uh.mask_prices(data, pit_intervals)

    stock_tickers = [
        t for t in data.columns
        if t not in etf_extra and data[t].notna().any()
    ]
    print(f"   Period: {data.index[0].date()} → {data.index[-1].date()}")
    print(f"   Tradable stocks: {len(stock_tickers)}")

    # --- Run all strategies ---
    print("\n[2/3] Running strategies...")
    results: dict[str, pd.Series] = {}

    def run(name, strategy, price_data):
        print(f"   {name}...")
        eq = backtest(strategy, price_data, initial_capital=CAPITAL,
                      transaction_cost=TX_COST, fixed_fee=FIXED_FEE)
        results[name] = eq

    # 1-2. Regime + momentum top-N
    for n in (10, 20):
        run(f"Regime+Mom Top{n}",
            RegimeStockPicker(stock_tickers=stock_tickers, top_n=n),
            data)

    # 3. Regime + recovery-momentum top-10
    run("Regime+RecMom Top10",
        RegimeRecoveryPicker(stock_tickers=stock_tickers, top_n=10),
        data)

    # 4. Regime + recovery-momentum top-20
    run("Regime+RecMom Top20",
        RegimeRecoveryPicker(stock_tickers=stock_tickers, top_n=20),
        data)

    # 5. Pure recovery momentum (no regime) — baseline
    run("RecoveryMom Top10 (pure)",
        RecoveryMomentumStrategy(top_n=10),
        data[stock_tickers])

    run("RecoveryMom Top20 (pure)",
        RecoveryMomentumStrategy(top_n=20),
        data[stock_tickers])

    # 6. TrendRider V7 (leveraged ETFs)
    etf_cols = [t for t in ["SPY", "TQQQ", "UPRO", "GLD", "DBC", "SHY"]
                if t in data.columns]
    run("TrendRider V7 (3x ETFs)", TrendRiderV7(), data[etf_cols])

    # 7. SPY benchmark
    spy = data["SPY"].dropna()
    results["SPY (benchmark)"] = (spy / spy.iloc[0]) * CAPITAL

    # --- Report ---
    print(f"\n[3/3] Results ({YEARS}y, ${CAPITAL:,.0f}, tx={TX_COST*100:.1f}bps + ${FIXED_FEE:.0f}/trade)")
    print("=" * 100)
    hdr = (f"{'Strategy':<30} {'Ann%':>8} {'Vol%':>8} {'Sharpe':>8} "
           f"{'Sortino':>8} {'MaxDD%':>8} {'Calmar':>8} {'WinRate':>8}")
    print(hdr)
    print("-" * 100)

    for name, eq in results.items():
        m = metrics.raw_summary(eq)
        print(f"{name:<30} {m['annualizedReturn']*100:>7.1f}% "
              f"{m['annualizedVolatility']*100:>7.1f}% "
              f"{m['sharpeRatio']:>8.2f} {m['sortinoRatio']:>8.2f} "
              f"{m['maxDrawdown']*100:>7.1f}% {m['calmarRatio']:>8.2f} "
              f"{m['winRate']*100:>7.1f}%")

    print("=" * 100)

    # Yearly breakdown for top strategies
    print("\n--- Yearly Returns ---")
    yearly: dict[str, dict[str, float]] = {}
    for name, eq in results.items():
        yr = {}
        for year, grp in eq.groupby(eq.index.year):
            if len(grp) >= 2:
                yr[str(year)] = grp.iloc[-1] / grp.iloc[0] - 1
        yearly[name] = yr

    all_years = sorted(set(y for d in yearly.values() for y in d))
    header = f"{'Year':<6}" + "".join(f"{name[:20]:>22}" for name in results)
    print(header)
    print("-" * len(header))
    for year in all_years:
        cols = []
        for name in results:
            r = yearly[name].get(year)
            cols.append(f"{r*100:>20.1f}%" if r is not None else f"{'—':>21}")
        print(f"{year:<6}" + "".join(cols))


if __name__ == "__main__":
    main()