"""
Round 2: Strategy improvement iteration.

Tests Hybrid Alpha variants that combine FactorCombo signal with inv-vol weighting,
and RecoveryQualityBlend that uses all strong factors without restrictive gates.
"""

import numpy as np
import pandas as pd

import data_manager
from universe import UNIVERSES
from main import backtest

# Top performers from round 1
from strategies.recovery_momentum import RecoveryMomentumStrategy
from strategies.factor_combo import FactorComboStrategy
from strategies.improved_momentum_quality import ImprovedMomentumQualityStrategy

# Round 2 strategies
from strategies.hybrid_alpha import HybridAlphaStrategy, RecoveryQualityBlendStrategy


def annual_return(eq: pd.Series) -> float:
    return eq.iloc[-1] / eq.iloc[0] - 1

def max_dd(eq: pd.Series) -> float:
    return ((eq / eq.cummax()) - 1).min()

def sharpe(eq: pd.Series) -> float:
    daily = eq.pct_change().dropna()
    if daily.std() == 0:
        return 0.0
    return (daily.mean() * 252) / (daily.std() * np.sqrt(252))

def sortino(eq: pd.Series) -> float:
    daily = eq.pct_change().dropna()
    downside = daily[daily < 0].std() * np.sqrt(252)
    if downside == 0:
        return 0.0
    return (daily.mean() * 252) / downside

def cagr(eq: pd.Series) -> float:
    yrs = (eq.index[-1] - eq.index[0]).days / 365.25
    if yrs <= 0:
        return 0.0
    return (eq.iloc[-1] / eq.iloc[0]) ** (1 / yrs) - 1

def calmar(eq: pd.Series) -> float:
    dd = max_dd(eq)
    if dd >= 0:
        return 0.0
    return cagr(eq) / abs(dd)


def main():
    universe = UNIVERSES["us"]
    tickers = universe["fetch"]()
    benchmark = universe["benchmark"]
    all_tickers = sorted(set(tickers + [benchmark]))

    data = data_manager.update("us", all_tickers, with_open=False)
    tickers = [t for t in tickers if t in data.columns]
    top_n = max(5, len(tickers) // 10)

    print(f"Universe: {len(tickers)} stocks + {benchmark}. top_n={top_n}")
    print(f"Data range: {data.index[0].date()} to {data.index[-1].date()}")

    strategies = {
        # === BASELINES (top 3 from round 1) ===
        "Recovery+Mom Top20 (base)": (
            RecoveryMomentumStrategy(top_n=20),
            data[tickers]
        ),
        "FactorCombo rec+deep (base)": (
            FactorComboStrategy(signal_name="rec_mfilt+deep_upvol", rebal_freq=21, top_n=20),
            data[tickers]
        ),
        "Improved MomQuality (base)": (
            ImprovedMomentumQualityStrategy(top_n=20),
            data[tickers]
        ),

        # === ROUND 2: HYBRID ALPHA ===
        "Hybrid InvVol Top20": (
            HybridAlphaStrategy(rebal_freq=21, top_n=20, use_invvol=True, regime_dampen=1.0),
            data[tickers]
        ),
        "Hybrid InvVol Top30": (
            HybridAlphaStrategy(rebal_freq=21, top_n=30, use_invvol=True, regime_dampen=1.0),
            data[tickers]
        ),
        "Hybrid EW Top20": (
            HybridAlphaStrategy(rebal_freq=21, top_n=20, use_invvol=False, regime_dampen=1.0),
            data[tickers]
        ),
        "Hybrid InvVol Dampen": (
            HybridAlphaStrategy(rebal_freq=21, top_n=20, use_invvol=True, regime_dampen=0.5),
            data[tickers]
        ),
        "Hybrid Biweekly": (
            HybridAlphaStrategy(rebal_freq=10, top_n=20, use_invvol=True, regime_dampen=1.0),
            data[tickers]
        ),

        # === ROUND 2: RECOVERY QUALITY BLEND ===
        "RecQuality Blend Top20": (
            RecoveryQualityBlendStrategy(top_n=20, rebal_freq=21),
            data[tickers]
        ),
        "RecQuality Blend Top30": (
            RecoveryQualityBlendStrategy(top_n=30, rebal_freq=21),
            data[tickers]
        ),
        "RecQuality Blend Biweekly": (
            RecoveryQualityBlendStrategy(top_n=20, rebal_freq=10),
            data[tickers]
        ),
    }

    # Run backtests
    equity = {}
    for name, (strat, strat_data) in strategies.items():
        print(f"Running {name}...")
        equity[name] = backtest(strat, strat_data, initial_capital=10_000)

    # SPY benchmark
    bench = data[benchmark].dropna()
    equity["SPY"] = (bench / bench.iloc[0]) * 10_000

    eq_df = pd.DataFrame(equity).sort_index()

    # Yearly returns
    years = list(range(2016, 2027))
    rows = []
    for yr in years:
        start = pd.Timestamp(f"{yr}-01-01")
        end = pd.Timestamp(f"{yr}-12-31")
        window = eq_df.loc[(eq_df.index >= start) & (eq_df.index <= end)].dropna(how="all")
        if window.empty:
            continue
        row = {"Year": yr}
        for col in eq_df.columns:
            s = window[col].dropna()
            if len(s) < 2:
                row[col] = np.nan
            else:
                row[col] = annual_return(s)
        rows.append(row)

    yr_df = pd.DataFrame(rows).set_index("Year")

    print("\n" + "=" * 80)
    print("YEARLY TOTAL RETURN (%)")
    print("=" * 80)
    print((yr_df * 100).round(2).to_string())

    # Excess over SPY
    excess = yr_df.sub(yr_df["SPY"], axis=0).drop(columns=["SPY"])
    print("\n" + "=" * 80)
    print("EXCESS vs SPY (pp)")
    print("=" * 80)
    print((excess * 100).round(2).to_string())

    # Full-period summary
    print("\n" + "=" * 80)
    print("FULL-PERIOD METRICS (sorted by Calmar)")
    print("=" * 80)

    summary_rows = []
    for col in eq_df.columns:
        eq = eq_df[col].dropna()
        if len(eq) < 252:
            continue
        summary_rows.append({
            "Strategy": col,
            "CAGR %": cagr(eq) * 100,
            "Sharpe": sharpe(eq),
            "Sortino": sortino(eq),
            "Max DD %": max_dd(eq) * 100,
            "Calmar": calmar(eq),
            "Win vs SPY": f"{(excess[col] > 0).sum()}/{len(excess)}" if col in excess.columns else "-",
        })

    summary = pd.DataFrame(summary_rows).sort_values("Calmar", ascending=False)
    pd.set_option('display.max_columns', None)
    pd.set_option('display.width', 200)
    print(summary.to_string(index=False))

    # Turnover analysis
    print("\n" + "=" * 80)
    print("TURNOVER ANALYSIS")
    print("=" * 80)
    for name, (strat, strat_data) in strategies.items():
        w = strat.generate_signals(strat_data)
        avg_turn = w.diff().abs().sum(axis=1).mean()
        print(f"  {name:<35s} avg daily turnover: {avg_turn:.4f}")


if __name__ == "__main__":
    main()