"""
Round 3: Signal-level ensemble and enhanced factor combo.

Focus: improve on FactorCombo's 34.6% CAGR / 1.02 Calmar by:
1. Ensembling two best signals for pick diversification
2. Adding momentum as a tiebreaker signal
3. Concentrating in fewer high-conviction names
4. Tail-risk protection only in extreme drawdowns
"""

import numpy as np
import pandas as pd

import data_manager
from universe import UNIVERSES
from main import backtest

from strategies.recovery_momentum import RecoveryMomentumStrategy
from strategies.factor_combo import FactorComboStrategy
from strategies.improved_momentum_quality import ImprovedMomentumQualityStrategy
from strategies.ensemble_alpha import EnsembleAlphaStrategy, EnhancedFactorComboStrategy


def annual_return(eq): return eq.iloc[-1] / eq.iloc[0] - 1
def max_dd(eq): return ((eq / eq.cummax()) - 1).min()
def sharpe(eq):
    d = eq.pct_change().dropna()
    return (d.mean() * 252) / (d.std() * np.sqrt(252)) if d.std() > 0 else 0
def sortino(eq):
    d = eq.pct_change().dropna()
    ds = d[d < 0].std() * np.sqrt(252)
    return (d.mean() * 252) / ds if ds > 0 else 0
def cagr(eq):
    yrs = (eq.index[-1] - eq.index[0]).days / 365.25
    return (eq.iloc[-1] / eq.iloc[0]) ** (1 / yrs) - 1 if yrs > 0 else 0
def calmar(eq):
    dd = max_dd(eq)
    return cagr(eq) / abs(dd) if dd < 0 else 0


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]

    print(f"Universe: {len(tickers)} stocks, data: {data.index[0].date()} to {data.index[-1].date()}")

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

        # Round 3: Ensemble
        "Ensemble Top20": (
            EnsembleAlphaStrategy(top_n=20, tail_protection=False),
            data[tickers]
        ),
        "Ensemble Top15": (
            EnsembleAlphaStrategy(top_n=15, tail_protection=False),
            data[tickers]
        ),
        "Ensemble Top20 +Tail": (
            EnsembleAlphaStrategy(top_n=20, tail_protection=True, tail_threshold=-0.15, tail_scale=0.5),
            data[tickers]
        ),
        "Ensemble Top20 +Tail10": (
            EnsembleAlphaStrategy(top_n=20, tail_protection=True, tail_threshold=-0.10, tail_scale=0.5),
            data[tickers]
        ),

        # Round 3: Enhanced FactorCombo
        "EnhFC Top15 mom20%": (
            EnhancedFactorComboStrategy(top_n=15, mom_boost=0.2, tail_protection=False),
            data[tickers]
        ),
        "EnhFC Top20 mom20%": (
            EnhancedFactorComboStrategy(top_n=20, mom_boost=0.2, tail_protection=False),
            data[tickers]
        ),
        "EnhFC Top15 mom30%": (
            EnhancedFactorComboStrategy(top_n=15, mom_boost=0.3, tail_protection=False),
            data[tickers]
        ),
        "EnhFC Top20 +Tail": (
            EnhancedFactorComboStrategy(top_n=20, mom_boost=0.2, tail_protection=True),
            data[tickers]
        ),
        "EnhFC Top10 mom20%": (
            EnhancedFactorComboStrategy(top_n=10, mom_boost=0.2, tail_protection=False),
            data[tickers]
        ),
    }

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

    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:
        window = eq_df.loc[f"{yr}"].dropna(how="all") if f"{yr}" in eq_df.index.strftime("%Y").unique() else pd.DataFrame()
        if window.empty:
            continue
        row = {"Year": yr}
        for col in eq_df.columns:
            s = window[col].dropna()
            row[col] = annual_return(s) if len(s) >= 2 else np.nan
        rows.append(row)

    yr_df = pd.DataFrame(rows).set_index("Year")
    excess = yr_df.sub(yr_df["SPY"], axis=0).drop(columns=["SPY"])

    print("\n" + "=" * 100)
    print("YEARLY RETURNS (%)")
    print("=" * 100)
    print((yr_df * 100).round(1).to_string())

    print("\n" + "=" * 100)
    print("FULL-PERIOD METRICS")
    print("=" * 100)
    print(f"{'Strategy':<30s} {'CAGR%':>7s} {'Sharpe':>7s} {'Sortino':>8s} {'MaxDD%':>8s} {'Calmar':>7s} {'WinSPY':>7s}")
    print("-" * 78)

    results = []
    for col in eq_df.columns:
        eq = eq_df[col].dropna()
        if len(eq) < 252:
            continue
        wins = (excess[col] > 0).sum() if col in excess.columns else 0
        total = len(excess) if col in excess.columns else 0
        results.append((col, cagr(eq)*100, sharpe(eq), sortino(eq), max_dd(eq)*100, calmar(eq), f"{wins}/{total}"))

    results.sort(key=lambda x: -x[5])  # sort by Calmar
    for r in results:
        print(f"{r[0]:<30s} {r[1]:>7.1f} {r[2]:>7.2f} {r[3]:>8.2f} {r[4]:>8.1f} {r[5]:>7.2f} {r[6]:>7s}")


if __name__ == "__main__":
    main()