"""
Export phase-averaged equity curves of the final candidates and baselines for
the report (weekly samples), plus the US/CN correlation of the new strategies.

    PYTHONPATH=. uv run python -m research.lab.export_curves
"""

import json
import os
import warnings

import numpy as np
import pandas as pd

from research.lab.evaluate import phase_runs, score_for
from research.lab.factor_scan import RESULTS, load
from research.lab.report import stats

warnings.filterwarnings("ignore")

SETS = {
    "us": ("2006-09-27", "2026-09-24", {
        "新: overnight+recovery": ("overnight_mom+rec_126", dict(top_n=10, every=10, buffer=2.0)),
        "现: rec_mfilt+deep_upvol 月频": ("baseline:rec_mfilt+deep_upvol", dict(top_n=10, every=21)),
        "现: mom7m+rec126 月频": ("baseline:mom7m+rec126", dict(top_n=10, every=21)),
        "等权 S&P 500 (PIT)": ("ew", dict(top_n=10000, every=5)),
    }),
    "cn_hs300": ("2010-01-04", "2026-09-24", {
        "新: 低风险动量4因子": ("amp_cut_mom+low_beta+pv_corr+sharpe_mom", dict(top_n=15, every=10, buffer=2.0)),
        "新: low_beta+sharpe_mom": ("low_beta+sharpe_mom", dict(top_n=20, every=10, buffer=1.0)),
        "现: down_resil+qual_mom 双周": ("baseline:down_resil+qual_mom", dict(top_n=10, every=10)),
        "现: up_cap+quality_mom 周频": ("baseline:up_cap+quality_mom", dict(top_n=10, every=5)),
        "等权 CSI 300 (PIT)": ("ew", dict(top_n=10000, every=5)),
    }),
}


def main() -> None:
    out, daily = {}, {}
    for market, (a, b, cands) in SETS.items():
        p, _ = load(market)
        curves = {}
        for name, (combo, kw) in cands.items():
            s = score_for(p, combo, market)
            runs = phase_runs(p, s, a, b, **kw)
            rets = pd.concat([r.returns.loc[a:b] for r in runs], axis=1).mean(axis=1)
            daily[(market, name)] = rets
            eq = (1 + rets).cumprod()
            wk = eq.resample("W-FRI").last().dropna()
            dd = eq / eq.cummax() - 1
            curves[name] = {"dates": [d.strftime("%Y-%m-%d") for d in wk.index],
                            "equity": [round(float(v), 4) for v in wk.values],
                            "drawdown": [round(float(v), 4) for v in dd.resample("W-FRI").min().dropna().values],
                            "stats": {k: float(v) for k, v in stats(eq).items()}}
            print(market, name, {k: round(v, 3) for k, v in curves[name]["stats"].items()}, flush=True)
        out[market] = curves
    us = daily[("us", "新: overnight+recovery")]
    cn = daily[("cn_hs300", "新: 低风险动量4因子")]
    # compound each market on its own calendar before joining: the exchanges
    # have different holidays, and an inner join on days would drop returns
    m = pd.concat([us, cn], axis=1, keys=["us", "cn"]).dropna().loc["2016-09-27":]
    mon = pd.concat({k: (1 + v.loc["2016-09-27":]).resample("ME").prod() - 1
                     for k, v in (("us", us), ("cn", cn))}, axis=1).dropna()
    cross = {"daily_corr": float(m.corr().iloc[0, 1]), "monthly_corr": float(mon.corr().iloc[0, 1])}
    for col in ["us", "cn", "5050"]:
        r = mon.mean(axis=1) if col == "5050" else mon[col]
        eq = (1 + r).cumprod()
        cross[col] = {"cagr": float(eq.iloc[-1] ** (12 / len(r)) - 1),
                      "sharpe_monthly": float(r.mean() / r.std() * np.sqrt(12)),
                      "mdd_monthly": float((eq / eq.cummax() - 1).min())}
    out["cross"] = cross
    print("cross", out["cross"])
    with open(os.path.join(RESULTS, "curves.json"), "w") as fh:
        json.dump(out, fh, ensure_ascii=False, default=float)


if __name__ == "__main__":
    main()