"""
Forward stepwise selection of equal-weight rank composites, scored on the
DESIGN period only; historical validation is reported separately. Previously inspected history
is not a fresh out-of-sample test.

    PYTHONPATH=. uv run python -m research.lab.stepwise --market us --max-k 4

Every evaluated composite is appended to the trial ledger
(data/research/results/trials_{market}.jsonl) so deflated Sharpe can count them.
"""

import argparse
import os
import time
import warnings

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS, load, phase_avg
from research.lab.factors import FACTORS, zrank
from research.lab.trials import append_trials

warnings.filterwarnings("ignore")

# Four ~5-year blocks; "maximin" scores a composite by its worst block Sharpe.
SUBPERIODS = {
    "us": [("2006-09-27", "2011-09-26"), ("2011-09-27", "2016-09-26"),
           ("2016-09-27", "2021-09-26"), ("2021-09-27", "2026-09-24")],
    "cn": [("2010-01-01", "2013-06-30"), ("2013-07-01", "2016-09-26"),
           ("2016-09-27", "2021-09-26"), ("2021-09-27", "2026-09-24")],
}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--market", default="us")
    ap.add_argument("--max-k", type=int, default=4)
    ap.add_argument("--top-n", type=int, default=20)
    ap.add_argument("--every", type=int, default=5)
    ap.add_argument("--beam", type=int, default=3)
    ap.add_argument("--objective", default="sharpe")
    ap.add_argument("--exclude", nargs="*", default=[])
    ap.add_argument("--buffer", type=float, default=0.0)
    args = ap.parse_args()
    mk = "us" if args.market == "us" else "cn"
    p, periods = load(args.market)
    names = [n for n, f in FACTORS.items() if mk in f.markets and n not in args.exclude]
    t = time.time()
    ranks = {}
    for n in names:
        try:
            ranks[n] = zrank(FACTORS[n](p), p.member)
        except Exception as exc:  # noqa: BLE001
            print(f"skip {n}: {exc}")
    names = list(ranks)
    print(f"{len(names)} factor ranks in {time.time() - t:.0f}s", flush=True)
    a, b = periods["design"]
    ledger = []
    subs = [(max(sa, a), min(sb, b)) for sa, sb in SUBPERIODS[mk]
            if max(sa, a) <= min(sb, b)]

    def evaluate(combo: tuple[str, ...]) -> dict:
        s = sum(ranks[n] for n in combo) / len(combo)
        if args.objective == "maximin":
            per = [phase_avg(p, s, sa, sb, args.top_n, args.every, buffer=args.buffer) for sa, sb in subs]
            st = {k: float(np.mean([x[k] for x in per])) for k in per[0]}
            st["maximin"] = min(x["sharpe"] for x in per)
            st["sub_sharpes"] = " ".join(f"{x['sharpe']:.2f}" for x in per)
            st["sub_cagrs"] = " ".join(f"{x['cagr']*100:.0f}" for x in per)
        else:
            st = phase_avg(p, s, a, b, args.top_n, args.every, buffer=args.buffer)
        ledger.append({"combo": "+".join(combo), "objective": args.objective, **st})
        return st

    beam = [()]
    best_overall = []
    for k in range(1, args.max_k + 1):
        cands = {}
        for base in beam:
            for n in names:
                if n in base:
                    continue
                combo = tuple(sorted(base + (n,)))
                if combo in cands:
                    continue
                cands[combo] = evaluate(combo)
        ranked = sorted(cands.items(), key=lambda kv: -kv[1][args.objective])
        beam = [c for c, _ in ranked[:args.beam]]
        print(f"--- k={k}: {len(cands)} combos", flush=True)
        for c, st in ranked[:6]:
            extra = (f" min {st['maximin']:.2f} subSh [{st['sub_sharpes']}] subCAGR [{st['sub_cagrs']}]"
                     if "maximin" in st else "")
            print(f"  {'+'.join(c):55s} sh {st['sharpe']:.2f} cagr {st['cagr']*100:5.1f}% "
                  f"mdd {st['mdd']*100:5.1f}% turn {st['turnover']:.0f}x{extra}", flush=True)
        best_overall += ranked[:args.beam]

    if args.objective == "maximin":
        _save(ledger, args.market, {"selection_start": a, "selection_end": b,
                                    "role": "development", "config": vars(args)})
        return
    # holdout, once, for the beam winners of each k
    print("=== historical validation of design winners (previously inspected data) ===")
    ha, hb = periods["holdout"]
    for c, st in best_overall:
        s = sum(ranks[n] for n in c) / len(c)
        hs = phase_avg(p, s, ha, hb, args.top_n, args.every)
        print(f"  {'+'.join(c):55s} D sh {st['sharpe']:.2f} cagr {st['cagr']*100:5.1f}% | "
              f"H sh {hs['sharpe']:.2f} cagr {hs['cagr']*100:5.1f}% mdd {hs['mdd']*100:5.1f}%", flush=True)
    _save(ledger, args.market, {"selection_start": a, "selection_end": b,
                                "role": "development", "config": vars(args)})


def _save(ledger: list[dict], market: str, protocol: dict) -> None:
    os.makedirs(RESULTS, exist_ok=True)
    path = os.path.join(RESULTS, f"trials_{market}.jsonl")
    append_trials(path, ledger, protocol)


if __name__ == "__main__":
    main()