"""
Round 2: Risk-Managed Ensemble with DD-reactive approach.

Key insight from R1: vol-target uniformly compresses returns (including uptrends),
losing too much CAGR. New approach: only cut exposure DURING drawdowns, not globally.
"""

import os
import sys

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import data_manager
from universe import UNIVERSES
from main import backtest
from strategies.ensemble_alpha import (
    EnsembleAlphaStrategy,
    RiskManagedEnsembleStrategy,
)


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 realized_vol(eq):
    return eq.pct_change().dropna().std() * np.sqrt(252)


def block_bootstrap(returns, n_boot=5000, block_len=21, seed=42):
    r = returns.values
    n = len(r)
    rng = np.random.default_rng(seed)
    n_blocks = int(np.ceil(n / block_len))
    span_years = n / 252.0
    cagrs = np.empty(n_boot)
    sharpes = np.empty(n_boot)
    mdds = np.empty(n_boot)
    for b in range(n_boot):
        starts = rng.integers(0, n - block_len + 1, size=n_blocks)
        idx = (starts[:, None] + np.arange(block_len)[None, :]).ravel()[:n]
        sample = r[idx]
        equity = np.cumprod(1.0 + sample)
        cagrs[b] = equity[-1] ** (1.0 / span_years) - 1.0
        std = sample.std(ddof=1)
        sharpes[b] = (sample.mean() / std * np.sqrt(252)) if std > 0 else 0.0
        running_max = np.maximum.accumulate(equity)
        mdds[b] = float(np.min(equity / running_max - 1.0))
    return pd.DataFrame({"cagr": cagrs, "sharpe": sharpes, "max_drawdown": mdds})


IS_END = "2022-12-31"
OOS_START = "2023-01-01"


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]
    stock_data = data[tickers]

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

    # =========================================================================
    # Baseline
    # =========================================================================
    base = EnsembleAlphaStrategy(top_n=10, tail_protection=False)
    eq_base = backtest(base, stock_data, initial_capital=10_000)

    print(f"\nBaseline (no RM):  CAGR={cagr(eq_base)*100:.1f}% Sharpe={sharpe(eq_base):.2f} MaxDD={max_dd(eq_base)*100:.1f}% Vol={realized_vol(eq_base)*100:.1f}%")

    # =========================================================================
    # Parameter sweep: DD-reactive approach
    # =========================================================================
    print("\n" + "=" * 110)
    print("  DD-REACTIVE RISK MANAGEMENT SWEEP")
    print("=" * 110)
    print(f"  {'Config':<55s} {'CAGR%':>7s} {'Sharpe':>7s} {'Sortino':>8s} {'MaxDD%':>8s} {'Calmar':>7s} {'Vol%':>6s}")
    print("  " + "-" * 98)

    configs = []
    for dd_fl in [0.15, 0.20, 0.25, 0.30, 0.40]:
        for dd_dn in [0.15, 0.20, 0.25, 0.30]:
            for vsg in [True, False]:
                for vsf in [0.40, 0.50, 0.60] if vsg else [0.50]:
                    strat = RiskManagedEnsembleStrategy(
                        top_n=10,
                        dd_floor=dd_fl, dd_denom=dd_dn,
                        vol_spike_guard=vsg, vol_spike_floor=vsf,
                    )
                    eq = backtest(strat, stock_data, initial_capital=10_000)
                    label = f"fl={dd_fl:.2f} dn={dd_dn:.2f} vsg={'Y' if vsg else 'N'} vsf={vsf:.2f}"
                    c = cagr(eq); s = sharpe(eq); so = sortino(eq)
                    mdd = max_dd(eq); cal = calmar(eq); rv = realized_vol(eq)
                    configs.append({
                        "label": label, "dd_floor": dd_fl, "dd_denom": dd_dn,
                        "vsg": vsg, "vsf": vsf,
                        "CAGR": c, "Sharpe": s, "Sortino": so,
                        "MaxDD": mdd, "Calmar": cal, "Vol": rv, "equity": eq,
                    })
                    # Only print selected configs to keep output manageable
                    if dd_dn in [0.20, 0.25] and dd_fl in [0.20, 0.25, 0.30] and vsf in [0.50]:
                        print(f"  {label:<55s} {c*100:>7.1f} {s:>7.2f} {so:>8.2f} {mdd*100:>8.1f} {cal:>7.2f} {rv*100:>6.1f}")

    # =========================================================================
    # Find configs meeting targets
    # =========================================================================
    print("\n  --- MEETING CAGR>40%, Sharpe>1.5, MaxDD>-25% ---")
    meeting = [c for c in configs if c["CAGR"] > 0.40 and c["Sharpe"] > 1.5 and c["MaxDD"] > -0.25]
    if meeting:
        for c in sorted(meeting, key=lambda x: -x["Calmar"])[:8]:
            print(f"  ✓ {c['label']:<50s} CAGR={c['CAGR']*100:.1f}% Sharpe={c['Sharpe']:.2f} MaxDD={c['MaxDD']*100:.1f}% Calmar={c['Calmar']:.2f}")
    else:
        print("  (None)")
        # Relax criteria
        print("\n  --- MEETING CAGR>38%, Sharpe>1.4, MaxDD>-25% ---")
        meeting2 = [c for c in configs if c["CAGR"] > 0.38 and c["Sharpe"] > 1.4 and c["MaxDD"] > -0.25]
        if meeting2:
            for c in sorted(meeting2, key=lambda x: -x["Calmar"])[:8]:
                print(f"  → {c['label']:<50s} CAGR={c['CAGR']*100:.1f}% Sharpe={c['Sharpe']:.2f} MaxDD={c['MaxDD']*100:.1f}% Calmar={c['Calmar']:.2f}")

        print("\n  --- BEST CALMAR with CAGR>35% ---")
        hi = [c for c in configs if c["CAGR"] > 0.35]
        for c in sorted(hi, key=lambda x: -x["Calmar"])[:5]:
            print(f"  → {c['label']:<50s} CAGR={c['CAGR']*100:.1f}% Sharpe={c['Sharpe']:.2f} MaxDD={c['MaxDD']*100:.1f}% Calmar={c['Calmar']:.2f}")

        print("\n  --- BEST with MaxDD > -25% ---")
        lo_dd = [c for c in configs if c["MaxDD"] > -0.25]
        for c in sorted(lo_dd, key=lambda x: -x["CAGR"])[:5]:
            print(f"  → {c['label']:<50s} CAGR={c['CAGR']*100:.1f}% Sharpe={c['Sharpe']:.2f} MaxDD={c['MaxDD']*100:.1f}% Calmar={c['Calmar']:.2f}")

    # Pick best overall by Calmar with CAGR > 38%
    candidates = [c for c in configs if c["CAGR"] > 0.38]
    if not candidates:
        candidates = sorted(configs, key=lambda x: -x["Calmar"])
    best = max(candidates, key=lambda x: x["Calmar"])

    print(f"\n  >>> RECOMMENDED: {best['label']}")
    print(f"      CAGR={best['CAGR']*100:.1f}% Sharpe={best['Sharpe']:.2f} Sortino={best['Sortino']:.2f} MaxDD={best['MaxDD']*100:.1f}% Calmar={best['Calmar']:.2f} Vol={best['Vol']*100:.1f}%")

    # =========================================================================
    # IS/OOS for recommended
    # =========================================================================
    print("\n" + "=" * 110)
    print("  IS/OOS VALIDATION")
    print("=" * 110)

    rec_strat = RiskManagedEnsembleStrategy(
        top_n=10, dd_floor=best["dd_floor"], dd_denom=best["dd_denom"],
        vol_spike_guard=best["vsg"], vol_spike_floor=best["vsf"],
    )

    is_data = stock_data[stock_data.index <= IS_END]
    oos_data = stock_data[stock_data.index >= OOS_START]

    eq_is = backtest(rec_strat, is_data, initial_capital=10_000)
    eq_oos = backtest(rec_strat, oos_data, initial_capital=10_000)
    eq_base_is = backtest(base, is_data, initial_capital=10_000)
    eq_base_oos = backtest(base, oos_data, initial_capital=10_000)

    print(f"\n  {'Strategy':<25s} {'Window':<10s} {'CAGR%':>7s} {'Sharpe':>7s} {'MaxDD%':>8s} {'Calmar':>7s}")
    print("  " + "-" * 68)
    for nm, ei, eo in [("RiskManaged", eq_is, eq_oos), ("Base (no RM)", eq_base_is, eq_base_oos)]:
        print(f"  {nm:<25s} {'IS':<10s} {cagr(ei)*100:>7.1f} {sharpe(ei):>7.2f} {max_dd(ei)*100:>8.1f} {calmar(ei):>7.2f}")
        print(f"  {nm:<25s} {'OOS':<10s} {cagr(eo)*100:>7.1f} {sharpe(eo):>7.2f} {max_dd(eo)*100:>8.1f} {calmar(eo):>7.2f}")

    # =========================================================================
    # Bootstrap on recommended
    # =========================================================================
    print("\n" + "=" * 110)
    print("  BLOCK BOOTSTRAP (5000 resamples)")
    print("=" * 110)

    rets = best["equity"].pct_change().dropna()
    boot = block_bootstrap(rets)
    print(f"\n    P(CAGR > 40%)   = {(boot['cagr'] > 0.40).mean()*100:.1f}%")
    print(f"    P(CAGR > 30%)   = {(boot['cagr'] > 0.30).mean()*100:.1f}%")
    print(f"    P(Sharpe > 1.5) = {(boot['sharpe'] > 1.5).mean()*100:.1f}%")
    print(f"    P(Sharpe > 1.0) = {(boot['sharpe'] > 1.0).mean()*100:.1f}%")
    print(f"    P(MaxDD > -25%) = {(boot['max_drawdown'] > -0.25).mean()*100:.1f}%")
    print(f"    P(MaxDD > -30%) = {(boot['max_drawdown'] > -0.30).mean()*100:.1f}%")

    # =========================================================================
    # Yearly returns
    # =========================================================================
    print("\n" + "=" * 110)
    print("  YEARLY RETURNS")
    print("=" * 110)

    bench_eq = data[benchmark].dropna()
    bench_eq = (bench_eq / bench_eq.iloc[0]) * 10_000

    eq_df = pd.DataFrame({
        "Raw Ens10": eq_base,
        "RiskManaged": best["equity"],
        "SPY": bench_eq,
    }).sort_index()

    years = sorted(eq_df.index.year.unique())
    print(f"\n  {'Year':<6s} {'Raw%':>8s} {'RM%':>8s} {'SPY%':>8s} {'RM-SPY':>8s}")
    print("  " + "-" * 42)
    for yr in years:
        w = eq_df.loc[eq_df.index.year == yr].dropna(how="all")
        if w.empty or len(w) < 2:
            continue
        r_raw = annual_return(w["Raw Ens10"].dropna()) if len(w["Raw Ens10"].dropna()) >= 2 else 0
        r_rm = annual_return(w["RiskManaged"].dropna()) if len(w["RiskManaged"].dropna()) >= 2 else 0
        r_spy = annual_return(w["SPY"].dropna()) if len(w["SPY"].dropna()) >= 2 else 0
        print(f"  {yr:<6d} {r_raw*100:>8.1f} {r_rm*100:>8.1f} {r_spy*100:>8.1f} {(r_rm-r_spy)*100:>+8.1f}")

    # =========================================================================
    # Summary
    # =========================================================================
    print(f"\n{'='*110}")
    print(f"  FINAL: RiskManagedEnsembleStrategy")
    print(f"  Config: top_n=10, dd_floor={best['dd_floor']}, dd_denom={best['dd_denom']}, vsg={best['vsg']}, vsf={best['vsf']}")
    print(f"  CAGR={best['CAGR']*100:.1f}%  Sharpe={best['Sharpe']:.2f}  Sortino={best['Sortino']:.2f}  MaxDD={best['MaxDD']*100:.1f}%  Calmar={best['Calmar']:.2f}")
    print(f"  vs Raw: CAGR {(best['CAGR']-cagr(eq_base))*100:+.1f}pp  Sharpe {best['Sharpe']-sharpe(eq_base):+.2f}  MaxDD {(best['MaxDD']-max_dd(eq_base))*100:+.1f}pp")


if __name__ == "__main__":
    main()