"""
US construction and overlay sweep for a few fixed scores, under one set of
rules (PIT S&P 500 members, next-open fills, 10 bp costs, phase averaging).

    PYTHONPATH=. uv run --with numba python -m research.lab.us_sweep --part construction
    PYTHONPATH=. uv run --with numba python -m research.lab.us_sweep --part overlays

Every configuration is appended to data/research/results/trials_us_sweep.jsonl
so the deflated Sharpe of the final pick can count them.
"""

from __future__ import annotations

import argparse
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab import fundamentals as fu
from research.lab import overlays as ov
from research.lab.engine import Costs, simulate
from research.lab.factor_scan import RESULTS, load
from research.lab.report import stats
from research.lab.rotation import factor_ranks

warnings.filterwarnings("ignore")

# 2006-2026 includes the 2008 crash; fundamentals start in 2011
LONG = [("2006-09-27", "2011-09-26"), ("2011-09-27", "2016-09-26"),
        ("2016-09-27", "2021-09-26"), ("2021-09-27", "2026-09-24")]
SHORT = [("2011-01-01", "2014-12-31"), ("2015-01-01", "2018-12-31"),
         ("2019-01-01", "2022-12-31"), ("2023-01-01", "2026-09-24")]


def _blend(*dfs: pd.DataFrame) -> pd.DataFrame:
    stack = np.stack([d.values for d in dfs])
    with np.errstate(invalid="ignore"):
        m = np.nanmean(stack, axis=0)
    return pd.DataFrame(m, index=dfs[0].index, columns=dfs[0].columns)


def sector_rank(df: pd.DataFrame, p) -> pd.DataFrame:
    df = df.where(p.member)
    sec = p.extra["sector"]
    out = pd.DataFrame(np.nan, index=df.index, columns=df.columns)
    for s in sec.dropna().unique():
        cols = [c for c in sec.index[sec == s] if c in df.columns]
        out[cols] = df[cols].rank(axis=1, pct=True)
    return out


def scores(p) -> dict[str, tuple[pd.DataFrame, list]]:
    pr = factor_ranks(p, "us")
    F = fu.load(p)
    Q = _blend(*(sector_rank(F[k], p) for k in ("gp_a", "cfo_a", "roe", "accruals", "leverage")))
    mom = _blend(pr["mom_12_1"], pr["resid_mom"])
    low = _blend(pr["ivol"], pr["low_beta"])
    return {
        "champ": (_blend(pr["overnight_mom"], pr["rec_126"]).where(p.member), LONG),
        "mom_qual_low": (_blend(mom, Q, low).where(p.member), SHORT),
    }


def run(p, score, subs, exposure=None, **kw) -> dict:
    """Phase-averaged daily returns over the whole span plus block stats."""
    a, b = subs[0][0], subs[-1][1]
    every = kw.get("every", 10)
    w0 = p.dates.searchsorted(pd.Timestamp(a))
    sp = p.slice(p.dates[max(w0 - 5, 0)], b)
    rets = []
    for off in range(every):
        r = simulate(sp, score, offset=off, warmup=5, exposure=exposure, **kw)
        rets.append(r.returns.loc[a:b])
        if off == 0:
            turn = r.turnover.loc[a:b].sum() / ((pd.Timestamp(b) - pd.Timestamp(a)).days / 365.25)
    R = pd.concat(rets, axis=1)
    per_phase = [stats((1 + R[c]).cumprod()) for c in R.columns]
    avg = R.mean(axis=1)                      # = a book split into ``every`` tranches
    out = {"sharpe": np.mean([s["sharpe"] for s in per_phase]),
           "cagr": np.mean([s["cagr"] for s in per_phase]),
           "mdd": np.mean([s["mdd"] for s in per_phase]),
           "sharpe_min_phase": min(s["sharpe"] for s in per_phase),
           "tranche_sharpe": stats((1 + avg).cumprod())["sharpe"], "turn": turn,
           "_rets": avg}
    blocks = []
    for sa, sb in subs:
        blocks.append(np.mean([stats((1 + R[c].loc[sa:sb]).cumprod())["sharpe"] for c in R.columns]))
    out["blocks"] = blocks
    out["min_block"] = min(blocks)
    return out


def fmt(name, st) -> str:
    return (f"{name:48s} sh {st['sharpe']:.2f} [{st['sharpe_min_phase']:.2f}] tr {st['tranche_sharpe']:.2f} "
            f"cagr {st['cagr']*100:5.1f}% mdd {st['mdd']*100:5.1f}% turn {st['turn']:4.1f}x "
            f"min {st['min_block']:.2f} [{' '.join(f'{x:.2f}' for x in st['blocks'])}]")


def log_trials(rows, part):
    from research.lab.trials import append_trials
    records = [{k: v for k, v in r.items() if not k.startswith("_") and k != "blocks"} for r in rows]
    append_trials(os.path.join(RESULTS, "trials_us_sweep.jsonl"), records,
                  {"part": part, "role": "development; historical data already examined"})


def construction(p, S) -> list[dict]:
    rows = []
    grid = []
    for top in (10, 20, 30):
        for buf in (0.0, 1.0, 2.0):
            grid.append((f"top{top} buf{buf:g} equal", dict(top_n=top, buffer=buf, weighting="equal")))
    for top in (10, 20, 30):
        grid.append((f"top{top} buf1 zvol", dict(top_n=top, buffer=1.0, weighting="zvol")))
        grid.append((f"top{top} buf1 invvol", dict(top_n=top, buffer=1.0, weighting="invvol")))
    grid.append(("dyn 10-30 thr0.80 buf1 equal", dict(top_n=30, buffer=1.0, dyn_n=(10, 0.80))))
    grid.append(("dyn 10-30 thr0.75 buf1 equal", dict(top_n=30, buffer=1.0, dyn_n=(10, 0.75))))
    grid.append(("top20 buf1 every5", dict(top_n=20, buffer=1.0, every=5)))
    grid.append(("top20 buf1 every21", dict(top_n=20, buffer=1.0, every=21)))
    for sname, (score, subs) in S.items():
        print(f"=== {sname}", flush=True)
        for label, kw in grid:
            kw = {"every": 10, **kw}
            st = run(p, score, subs, **kw)
            rows.append({"score": sname, "config": label, **st})
            print(fmt(label, st), flush=True)
    return rows


def overlays(p, S, base_kw) -> list[dict]:
    lvl = p.benchmark.ffill()
    ew = ov.ew_index(p)
    X = {
        "none": None,
        "spy derisk_vol": ov.derisk_vol(lvl),
        "spy faber10m": ov.faber_sma(lvl),
        "spy jump_model": ov.jump_model(lvl, p.cash_ret),
        "spy dm_state24m": ov.momentum_state(lvl),
        "ew breadth-free derisk_vol": ov.derisk_vol(ew),
    }
    rows = []
    for sname, (score, subs) in S.items():
        print(f"=== {sname} {base_kw}", flush=True)
        base = run(p, score, subs, **base_kw)
        # strategy-level de-risk-only vol target from its own trailing returns
        own = base["_rets"]
        own_lvl = (1 + own).cumprod()
        X_s = dict(X)
        X_s["own derisk_vol"] = ov.derisk_vol(own_lvl, min_hist=126).reindex(p.dates).fillna(1.0)
        for name, e in X_s.items():
            st = base if e is None else run(p, score, subs, exposure=e, **base_kw)
            rows.append({"score": sname, "config": f"overlay {name}", **st})
            print(fmt(f"overlay {name}", st), flush=True)
    return rows


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--part", default="construction")
    ap.add_argument("--top-n", type=int, default=20)
    ap.add_argument("--buffer", type=float, default=1.0)
    ap.add_argument("--weighting", default="equal")
    ap.add_argument("--fee", type=float, default=0.0, help="fixed fee per order on $10k")
    args = ap.parse_args()
    p, _ = load("us")
    if args.fee:
        p = p.slice()
        p.costs = Costs(p.costs.buy, p.costs.sell, None, args.fee, 10_000.0)
    S = scores(p)
    if args.part == "construction":
        rows = construction(p, S)
    else:
        rows = overlays(p, S, dict(top_n=args.top_n, buffer=args.buffer, every=10,
                                   weighting=args.weighting))
    log_trials(rows, args.part)
    with open(os.path.join(RESULTS, f"us_sweep_{args.part}.pkl"), "wb") as fh:
        pickle.dump(rows, fh)


if __name__ == "__main__":
    main()