"""
Rebalancing frequency comparison: daily (1d) vs weekly (5d) vs biweekly (10d) vs monthly (21d).
Shows yearly returns and max drawdown for each frequency, for all champion strategies.
"""

from __future__ import annotations
import warnings
import numpy as np
import pandas as pd
import data_manager
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

REBAL_CONFIGS = [
    ("daily", 1),
    ("weekly", 5),
    ("biweekly", 10),
    ("monthly", 21),
]


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, rebal=21, cost=0.001):
    w = strat(prices, func, top_n=10, rebal=rebal)
    eq = bt(w, prices, cost=cost)
    return eq / eq.iloc[0] * INITIAL


def year_returns(eq: pd.Series) -> dict[int, float]:
    dr = eq.pct_change().fillna(0)
    return {y: float((1 + dr[dr.index.year == y]).prod() - 1)
            for y in sorted(dr.index.year.unique())}


def max_drawdown(eq: pd.Series) -> float:
    rm = eq.cummax()
    dd = (eq - rm) / rm
    return float(dd.min())


def max_drawdown_yearly(eq: pd.Series) -> dict[int, float]:
    result = {}
    for y in sorted(eq.index.year.unique()):
        chunk = eq[eq.index.year == y]
        if len(chunk) < 5:
            continue
        rm = chunk.cummax()
        dd = (chunk - rm) / rm
        result[y] = float(dd.min())
    return result


def cagr(eq: pd.Series) -> float:
    dr = eq.pct_change().dropna()
    if len(dr) < 100:
        return np.nan
    ny = len(dr) / 252
    tot = eq.iloc[-1] / eq.iloc[0] - 1
    return (1 + tot) ** (1 / ny) - 1


def sharpe(eq: pd.Series) -> float:
    dr = eq.pct_change().dropna()
    if len(dr) < 100 or dr.std() == 0:
        return np.nan
    return float(dr.mean() / dr.std() * np.sqrt(252))


def turnover_annual(func, prices, rebal):
    """Estimate annualised turnover (one-way)."""
    w = strat(prices, func, top_n=10, rebal=rebal)
    daily_turn = w.diff().abs().sum(axis=1).mean()
    return daily_turn * 252


def print_by_year(strat_defs, prices, bench_eq, bench_label, market_label, years):
    """For each year, print a table: strategies as rows, rebal frequencies as columns."""

    freq_labels = [r for r, _ in REBAL_CONFIGS]

    # Pre-compute all equities and returns
    all_eqs = {}  # {(sname, freq): equity}
    for sname, func in strat_defs.items():
        for rlabel, rdays in REBAL_CONFIGS:
            all_eqs[(sname, rlabel)] = run_equity(func, prices, rebal=rdays)

    all_rets = {}  # {(sname, freq): {year: ret}}
    for key, eq in all_eqs.items():
        all_rets[key] = year_returns(eq)

    bench_rets = year_returns(bench_eq)
    snames = list(strat_defs.keys())
    name_w = max(len(s) for s in snames) + 1

    for year in years:
        line_w = name_w + 4 + 20 * (len(freq_labels) + 1)
        print(f"\n{'=' * line_w}")
        print(f"  {market_label} — {year}  (fresh $10,000)")
        print(f"{'=' * line_w}")

        # Header
        print(f"  {'Strategy':<{name_w}}", end="")
        for f in freq_labels:
            print(f"  {f:>18}", end="")
        print(f"  {bench_label:>18}")
        print(f"  {'-'*name_w}", end="")
        for _ in range(len(freq_labels) + 1):
            print(f"  {'-'*18}", end="")
        print()

        for sname in snames:
            print(f"  {sname:<{name_w}}", end="")

            # Find best freq for this strategy this year
            freq_vals = {}
            for f in freq_labels:
                r = all_rets[(sname, f)].get(year)
                if r is not None and abs(r) > 0.0005:
                    freq_vals[f] = r

            best_f = max(freq_vals, key=freq_vals.get) if freq_vals else None

            for f in freq_labels:
                r = all_rets[(sname, f)].get(year)
                if r is not None and abs(r) > 0.0005:
                    v = INITIAL * (1 + r)
                    marker = " ★" if f == best_f else "  "
                    print(f"  ${v:>9,.0f} {r:>+5.0%}{marker}", end="")
                else:
                    print(f"  {'—':>18}", end="")

            # Benchmark (same for all strategies)
            br = bench_rets.get(year)
            if br is not None and abs(br) > 0.0005:
                print(f"  ${INITIAL*(1+br):>9,.0f} {br:>+5.0%}  ", end="")
            else:
                print(f"  {'—':>18}", end="")
            print()

        # Best strategy per freq
        print(f"  {'-'*name_w}", end="")
        for _ in range(len(freq_labels) + 1):
            print(f"  {'-'*18}", end="")
        print()

        print(f"  {'BEST':<{name_w}}", end="")
        for f in freq_labels:
            best_r = -999
            best_s = ""
            for sname in snames:
                r = all_rets[(sname, f)].get(year)
                if r is not None and abs(r) > 0.0005 and r > best_r:
                    best_r = r
                    best_s = sname
            if best_r > -999:
                print(f"  ${INITIAL*(1+best_r):>9,.0f} {best_r:>+5.0%}  ", end="")
            else:
                print(f"  {'—':>18}", end="")
        # bench
        br = bench_rets.get(year)
        if br is not None and abs(br) > 0.0005:
            print(f"  ${INITIAL*(1+br):>9,.0f} {br:>+5.0%}  ", end="")
        else:
            print(f"  {'—':>18}", end="")
        print()


def main():
    years = list(range(2015, 2027))

    # ===== US =====
    print(f"\n{'#'*130}")
    print(f"{'#'*50}  US MARKET  {'#'*50}")
    print(f"{'#'*130}")

    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,
    }

    print_by_year(us_strats, stocks_us, eq_spy, "SPY", "US", years)

    # ===== CN =====
    print(f"\n\n{'#'*130}")
    print(f"{'#'*50}  CN MARKET  {'#'*50}")
    print(f"{'#'*130}")

    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,
    }

    if bench_cn is not None:
        eq_csi = bench_cn / bench_cn.iloc[0] * INITIAL
    else:
        eq_csi = pd.Series(dtype=float)

    print_by_year(cn_strats, stocks_cn, eq_csi, "CSI300", "CN", years)


if __name__ == "__main__":
    main()