"""
Final evaluation of candidate strategies against the existing champions,
under identical rules (PIT universe, next-open execution, costs, phase
averaging). Writes data/research/results/final_{market}.json.

    PYTHONPATH=. uv run python -m research.lab.evaluate --market us
"""

import argparse
import json
import os
import warnings

import numpy as np
import pandas as pd

from research.lab import overlays as ov
from research.lab.construction import combo_score
from research.lab.engine import Costs, simulate
from research.lab.factor_scan import RESULTS, load
from research.lab.report import deflated_sharpe, stats, yearly
from research.lab.stepwise import SUBPERIODS
from research.lab.trials import count_trials

warnings.filterwarnings("ignore")

# name -> (combo, construction). "baseline:" combos are the current factor_combo champions.
CANDIDATES = {
    "us": {
        "champion rec_mfilt+deep_upvol monthly": ("baseline:rec_mfilt+deep_upvol", dict(top_n=10, every=21)),
        "champion mom7m+rec126 monthly": ("baseline:mom7m+rec126", dict(top_n=10, every=21)),
        "champion mom7m+rec126 weekly": ("baseline:mom7m+rec126", dict(top_n=10, every=5)),
        "EW S&P 500 (PIT)": ("ew", dict(top_n=10000, every=5)),
        "new overnight+rec": ("overnight_mom+rec_126", dict(top_n=10, every=10, buffer=2.0)),
        "new overnight6+rec": ("overnight_mom_6+rec_126", dict(top_n=10, every=10, buffer=2.0)),
        "new smooth-mom4": ("overnight_mom+overnight_mom_6+amp_cut_mom+rec_126", dict(top_n=10, every=10, buffer=2.0)),
        "new overnight+ampcut+rec": ("overnight_mom+amp_cut_mom+rec_126", dict(top_n=10, every=10, buffer=2.0)),
    },
    "cn": {
        "champion up_cap+quality_mom weekly": ("baseline:up_cap+quality_mom", dict(top_n=10, every=5)),
        "champion up_cap+mom_gap weekly": ("baseline:up_cap+mom_gap", dict(top_n=10, every=5)),
        "champion down_resil+qual_mom biweekly": ("baseline:down_resil+qual_mom", dict(top_n=10, every=10)),
        "EW CSI 300 (PIT)": ("ew", dict(top_n=10000, every=5)),
    },
}
LIVE_FEES = {"us": (2.0, 10_000.0), "cn": (5.0, 70_000.0)}
WINDOW = ("2016-09-27", "2026-09-24")          # the web UI's 10-year window


def score_for(p, combo, market):
    if combo == "ew":
        return p.close.notna().astype(float).where(p.member)
    return combo_score(p, combo, market)


def phase_runs(p, score, a, b, exposure=None, **kw):
    every = kw.get("every", 5)
    w0 = p.dates.searchsorted(pd.Timestamp(a))
    sp = p.slice(p.dates[max(w0 - 5, 0)], b)
    runs = []
    for off in range(every if kw.get("top_n", 10) < 1000 else 1):
        r = simulate(sp, score, offset=off, warmup=5, exposure=exposure, **kw)
        runs.append(r)
    return runs


def summarize(runs, a, b) -> dict:
    sts = [stats(r.equity.loc[a:b], r.turnover) for r in runs]
    df = pd.DataFrame(sts)
    rets = pd.concat([r.returns.loc[a:b] for r in runs], axis=1).mean(axis=1)
    eq = (1 + rets).cumprod()
    out = {k: float(df[k].mean()) for k in df.columns}
    out["cagr_phase_min"] = float(df["cagr"].min())
    out["cagr_phase_max"] = float(df["cagr"].max())
    out["yearly"] = {str(k.year): float(v) for k, v in yearly(eq).items()}
    out["_rets"] = rets
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--market", default="us")
    ap.add_argument("--add", nargs="*", default=[], help="extra 'name=combo|top|every|buffer'")
    ap.add_argument("--n-trials", type=int, default=0, help="override trial count for DSR")
    ap.add_argument("--registered", nargs="+", help="Validate registered strategies using the production ledger")
    ap.add_argument("--start")
    ap.add_argument("--end")
    ap.add_argument("--output")
    args = ap.parse_args()
    if args.registered:
        from research.lab.registered import evaluate_registered
        evaluate_registered(args)
        return
    mk = "us" if args.market == "us" else "cn"
    p, _ = load(args.market)
    cands = dict(CANDIDATES[mk])
    for spec in args.add:
        name, rest = spec.split("=", 1)
        combo, top, every, buf = rest.split("|")
        cands[name] = (combo, dict(top_n=int(top), every=int(every), buffer=float(buf)))
    n_trials = args.n_trials or count_trials(RESULTS, args.market)
    if n_trials < 1:
        raise ValueError("DSR needs a trial ledger or an explicit --n-trials count")
    full = (SUBPERIODS[mk][0][0], WINDOW[1])
    out = {}
    fee, cap = LIVE_FEES[mk]
    for name, (combo, kw) in cands.items():
        s = score_for(p, combo, args.market)
        res = {"combo": combo, "construction": kw}
        base_runs = phase_runs(p, s, *WINDOW, **kw)
        res["window"] = summarize(base_runs, *WINDOW)
        res["full"] = summarize(phase_runs(p, s, *full, **kw), *full)
        res["blocks"] = [summarize(phase_runs(p, s, a, b, **kw), a, b) for a, b in SUBPERIODS[mk]]
        # live-book fixed fee
        pf = p.slice(); pf.costs = Costs(p.costs.buy, p.costs.sell, p.costs.sell_tax, fee, cap)
        res["window_livefee"] = summarize(phase_runs(pf, s, *WINDOW, **kw), *WINDOW)
        # double proportional cost
        p2 = p.slice(); p2.costs = Costs(p.costs.buy * 2, p.costs.sell * 2, p.costs.sell_tax, 0.0, cap)
        res["window_2xcost"] = summarize(phase_runs(p2, s, *WINDOW, **kw), *WINDOW)
        res["window_exec_close"] = summarize(phase_runs(p, s, *WINDOW, exec_at="close", **kw), *WINDOW)
        # drawdown-control overlay: slow/fast market trend blend on the EW index
        ew = ov.ew_index(p)
        slow = (ew / ew.shift(252) > 1).astype(float)
        fast = (ew / ew.shift(21) > 1).astype(float)
        expo = (0.5 + 0.25 * slow + 0.25 * fast).where(slow + fast < 2, 1.0)
        res["window_trend_overlay"] = summarize(phase_runs(p, s, *WINDOW, exposure=expo, **kw), *WINDOW)
        dsr = deflated_sharpe(res["full"]["_rets"], n_trials)
        res["dsr_full"] = {**dsr, "observed_trials_lower_bound": n_trials,
                           "evidence_role": "development, not independent OOS"}
        out[name] = res
        w, f = res["window"], res["full"]
        print(f"{name:42s} 2016-26: CAGR {w['cagr']*100:5.1f}% [{w['cagr_phase_min']*100:4.1f},{w['cagr_phase_max']*100:4.1f}] "
              f"Sh {w['sharpe']:.2f} MDD {w['mdd']*100:5.1f}% | full: CAGR {f['cagr']*100:5.1f}% Sh {f['sharpe']:.2f} "
              f"MDD {f['mdd']*100:5.1f}% | blocks Sh {' '.join(f'{b_['sharpe']:.2f}' for b_ in res['blocks'])} "
              f"| livefee {res['window_livefee']['cagr']*100:5.1f}% 2x {res['window_2xcost']['cagr']*100:5.1f}% "
              f"close {res['window_exec_close']['cagr']*100:5.1f}% trend {res['window_trend_overlay']['cagr']*100:5.1f}%/"
              f"{res['window_trend_overlay']['sharpe']:.2f}/{res['window_trend_overlay']['mdd']*100:.0f}% "
              f"| DSR {dsr['dsr']:.2f} (N={n_trials})", flush=True)

    def clean(o):
        if isinstance(o, dict):
            return {k: clean(v) for k, v in o.items() if not k.startswith("_")}
        if isinstance(o, list):
            return [clean(v) for v in o]
        if isinstance(o, (np.floating, np.integer)):
            return float(o)
        return o
    os.makedirs(RESULTS, exist_ok=True)
    with open(os.path.join(RESULTS, f"final_{args.market}.json"), "w") as fh:
        json.dump(clean(out), fh, indent=1)


if __name__ == "__main__":
    main()