"""
Yearly ROI report for champion strategies vs SPY, starting from $10,000.
"""

from __future__ import annotations
import warnings
import numpy as np
import pandas as pd
import data_manager
from universe import UNIVERSES
from factor_loop import (
    strat, bt, stats, combo,
    f_rec_mom, f_rec_126, f_rec_63,
    f_mom_12_1, f_mom_6_1, f_mom_intermediate,
    f_above_ma200, f_golden_cross,
    f_up_volume_proxy, f_gap_up_freq,
    f_rec_mom_filtered, f_down_resilience,
    f_up_capture, f_52w_high, f_str_10d,
    f_earnings_drift, f_reversal_vol,
)

warnings.filterwarnings("ignore")

INITIAL = 10_000


def f_quality_mom(p):
    mom = f_mom_12_1(p)
    consist = (p.pct_change() > 0).astype(float).rolling(252, min_periods=126).mean()
    mom_r = mom.rank(axis=1, pct=True, na_option="keep")
    con_r = consist.rank(axis=1, pct=True, na_option="keep")
    up_r = f_up_volume_proxy(p).rank(axis=1, pct=True, na_option="keep")
    return 0.4 * mom_r + 0.3 * con_r + 0.3 * up_r


def f_mom_x_gap(p):
    return (f_mom_12_1(p).rank(axis=1, pct=True, na_option="keep") *
            f_gap_up_freq(p).rank(axis=1, pct=True, na_option="keep"))


def run_equity(func, prices, cost=0.001):
    w = strat(prices, func, top_n=10)
    eq = bt(w, prices, cost=cost)
    return eq / eq.iloc[0] * INITIAL


def yearly_table(equities: dict[str, pd.Series], title: str):
    print(f"\n{'='*130}")
    print(f"  {title}")
    print(f"  Starting capital: ${INITIAL:,.0f}")
    print(f"{'='*130}")

    names = list(equities.keys())
    all_years = sorted(set(y for eq in equities.values() for y in eq.index.year.unique()))

    # Header
    print(f"\n  {'Year':<6}", end="")
    for n in names:
        print(f" | {n[:24]:>24}", end="")
    print()
    print(f"  {'-'*6}", end="")
    for _ in names:
        print(f"-+-{'-'*24}", end="")
    print()

    # Track portfolio values
    year_end_vals = {n: INITIAL for n in names}

    for year in all_years:
        print(f"  {year:<6}", end="")
        for n in names:
            eq = equities[n]
            yr_data = eq[eq.index.year == year]
            if len(yr_data) < 2:
                print(f" | {'—':>24}", end="")
                continue
            start_val = yr_data.iloc[0]
            end_val = yr_data.iloc[-1]
            ret = end_val / start_val - 1
            year_end_vals[n] = end_val
            # Show both return % and portfolio value
            print(f" | {ret:>+7.1%}  ${end_val:>12,.0f}", end="")
        print()

    # Summary rows
    print(f"  {'-'*6}", end="")
    for _ in names:
        print(f"-+-{'-'*24}", end="")
    print()

    # Total return
    print(f"  {'Total':<6}", end="")
    for n in names:
        eq = equities[n]
        total = eq.iloc[-1] / INITIAL - 1
        print(f" | {total:>+7.0%}  ${eq.iloc[-1]:>12,.0f}", end="")
    print()

    # CAGR
    print(f"  {'CAGR':<6}", end="")
    for n in names:
        eq = equities[n]
        ny = len(eq) / 252
        total = eq.iloc[-1] / INITIAL - 1
        cagr = (1 + total) ** (1 / ny) - 1
        print(f" | {cagr:>+7.1%}  {'':>12}", end="")
    print()

    # Sharpe
    print(f"  {'Sharpe':<6}", end="")
    for n in names:
        eq = equities[n]
        dr = eq.pct_change().dropna()
        sh = dr.mean() / dr.std() * np.sqrt(252) if dr.std() > 0 else 0
        print(f" | {sh:>7.2f}  {'':>12}", end="")
    print()

    # Max DD
    print(f"  {'MaxDD':<6}", end="")
    for n in names:
        eq = equities[n]
        rm = eq.cummax()
        dd = ((eq - rm) / rm).min()
        print(f" | {dd:>+7.1%}  {'':>12}", end="")
    print()

    # Best/Worst year
    print(f"  {'Best':<6}", end="")
    for n in names:
        eq = equities[n]
        dr = eq.pct_change().fillna(0)
        yr_rets = {y: float((1 + dr[dr.index.year == y]).prod() - 1) for y in all_years}
        # skip warmup year
        active = {y: r for y, r in yr_rets.items() if abs(r) > 0.001}
        if active:
            best_y = max(active, key=active.get)
            print(f" | {active[best_y]:>+7.1%}  ({best_y})       ", end="")
        else:
            print(f" | {'—':>24}", end="")
    print()

    print(f"  {'Worst':<6}", end="")
    for n in names:
        eq = equities[n]
        dr = eq.pct_change().fillna(0)
        yr_rets = {y: float((1 + dr[dr.index.year == y]).prod() - 1) for y in all_years}
        active = {y: r for y, r in yr_rets.items() if abs(r) > 0.001}
        if active:
            worst_y = min(active, key=active.get)
            print(f" | {active[worst_y]:>+7.1%}  ({worst_y})       ", end="")
        else:
            print(f" | {'—':>24}", end="")
    print()


def main():
    # ===== US =====
    prices_us = data_manager.load("us")
    bench_us = prices_us["SPY"].dropna()
    stocks_us = prices_us.drop(columns=["SPY"], errors="ignore")

    eq_spy = bench_us / bench_us.iloc[0] * INITIAL

    us_strats = {
        "rec_mfilt+deep×upvol": combo([
            (f_rec_mom_filtered, 0.5),
            (combo([(f_rec_126, 0.5), (f_up_volume_proxy, 0.5)]), 0.5),
        ]),
        "ma200+mom7m+rec126": combo([
            (f_above_ma200, 0.33), (f_mom_intermediate, 0.33), (f_rec_126, 0.34)
        ]),
        "rec_mfilt+ma200": combo([
            (f_rec_mom_filtered, 0.5), (f_above_ma200, 0.5)
        ]),
        "mom7m+rec126": combo([
            (f_mom_intermediate, 0.5), (f_rec_126, 0.5)
        ]),
        "BASELINE:rec+mom": f_rec_mom,
    }

    us_equities = {}
    for name, func in us_strats.items():
        us_equities[name] = run_equity(func, stocks_us)
    us_equities["SPY (Benchmark)"] = eq_spy

    yearly_table(us_equities, "US MARKET — Champion Strategies vs SPY — $10,000 Starting Capital")

    # ===== CN =====
    prices_cn = data_manager.load("cn")
    bench_cn = prices_cn["000300.SS"].dropna() if "000300.SS" in prices_cn.columns else None
    stocks_cn = prices_cn.drop(columns=["000300.SS"], errors="ignore")

    cn_strats = {
        "up_cap+quality_mom": combo([
            (f_up_capture, 0.5), (f_quality_mom, 0.5)
        ]),
        "down_resil+qual_mom": combo([
            (f_down_resilience, 0.5), (f_quality_mom, 0.5)
        ]),
        "rec63+mom×gap": combo([
            (f_rec_63, 0.5), (f_mom_x_gap, 0.5)
        ]),
        "up_cap+mom×gap": combo([
            (f_up_capture, 0.5), (f_mom_x_gap, 0.5)
        ]),
        "BASELINE:rec+mom": f_rec_mom,
    }

    cn_equities = {}
    for name, func in cn_strats.items():
        cn_equities[name] = run_equity(func, stocks_cn)
    if bench_cn is not None:
        cn_equities["CSI300 (Benchmark)"] = bench_cn / bench_cn.iloc[0] * INITIAL

    yearly_table(cn_equities, "CN MARKET — Champion Strategies vs CSI 300 — $10,000 Starting Capital")


if __name__ == "__main__":
    main()