"""Mid/small factor scan with a survivorship experiment.

Compact top-N next-open backtests (full phase averaging, 20bp/side) on the
PIT S&P 400/600 union Panel, run against BOTH builds: the Yahoo-only panel
(pre-Tiingo backup, delisted names silently missing) and the Tiingo-extended
rebuild. A factor whose conclusion flips between the two was a survivorship
artifact; one that survives both is actionable.

Windows: design 2007-01..2016-12, holdout 2017-01..2026-09. Scores are
computed on the full history and the index is masked to PIT members, so
pre-admission bars inform indicators but never selection.

The overnight composite approximates the production robust_momentum sleeve
(overnight-return momentum + distance from the 126d low), not its exact code.

    PYTHONPATH=. uv run python -m research.midsmall_scan --panel both
"""
from __future__ import annotations

import argparse
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.engine import simulate
from research.lab.report import forward_returns, ic_series, ic_summary, stats

warnings.filterwarnings("ignore")

DESIGN = ("2007-01-01", "2016-12-31")
HOLDOUT = ("2017-01-01", "2026-09-24")
RESULTS = os.path.join("data", "research", "results")


def load(panel: str):
    if panel == "new":
        from research.lab.panel_mid import load_us_mid
        return load_us_mid("midsmall")
    with open(os.path.join("data", "research", "us_midsmall.panel.pre_tiingo.pkl"), "rb") as fh:
        return pickle.load(fh)


def factors(p) -> dict:
    c = p.close.where(p.member)
    o = p.open.where(p.member)
    vol = c.pct_change(fill_method=None).rolling(60, min_periods=40).std()
    onr = ((o / c.shift(1) - 1).where(lambda x: x.abs() < 0.5)
           .rolling(63, min_periods=40).mean())
    dist126 = c / c.rolling(126, min_periods=100).min() - 1
    hi252 = c / c.rolling(252, min_periods=200).max()
    pr = lambda x: x.rank(axis=1, pct=True)
    return {
        "mom_12_1": c.shift(21) / c.shift(252) - 1,
        "mom_6_1": c.shift(21) / c.shift(126) - 1,
        "rev_1m": -(c / c.shift(21) - 1),
        "rev_1w": -(c / c.shift(5) - 1),
        "lowvol": -vol,
        "hi252": hi252,
        "dist126": dist126,
        "onr63": onr,
        "orm_comp": pr(onr) + pr(dist126),
    }


def phase_avg(p, s, window, top_n, every) -> dict:
    a, b = window
    w0 = p.dates.searchsorted(pd.Timestamp(a))
    sp = p.slice(p.dates[max(w0 - 5, 0)], b)
    rows = []
    for off in range(every):
        r = simulate(sp, s, top_n=top_n, every=every, offset=off, warmup=5)
        rows.append(stats(r.equity.loc[a:b], r.turnover))
    return pd.DataFrame(rows).mean().to_dict()


def run(panel: str, top_n: int, every: int) -> pd.DataFrame:
    p = load(panel)
    fwd5 = forward_returns(p, 5)
    rows = []
    for name, s in factors(p).items():
        row = {"panel": panel, "factor": name, "top_n": top_n, "every": every}
        for wname, window in (("design", DESIGN), ("holdout", HOLDOUT)):
            st = phase_avg(p, s, window, top_n, every)
            ic = ic_summary(ic_series(s.loc[window[0]:window[1]], fwd5.loc[window[0]:window[1]],
                                      p.member.loc[window[0]:window[1]], 5))
            row.update({f"{wname}_cagr": st["cagr"], f"{wname}_sharpe": st["sharpe"],
                        f"{wname}_mdd": st["mdd"], f"{wname}_turn": st["turnover"],
                        f"{wname}_ic5": ic["ic"], f"{wname}_ic_t": ic["t"]})
        rows.append(row)
        print(f"{name:10s} D cagr {row['design_cagr']:6.1%} sh {row['design_sharpe']:5.2f} "
              f"ic {row['design_ic5']:+.3f} | H cagr {row['holdout_cagr']:6.1%} "
              f"sh {row['holdout_sharpe']:5.2f} ic {row['holdout_ic5']:+.3f}", flush=True)
    return pd.DataFrame(rows)


def ew_universe(p) -> None:
    r = p.close.pct_change(fill_method=None).where(p.member.shift(1, fill_value=False))
    ew = (1 + r.mean(axis=1).fillna(0.0)).cumprod()
    b = p.benchmark
    for wname, (a, bnd) in (("design", DESIGN), ("holdout", HOLDOUT)):
        se = stats(ew.loc[a:bnd])
        sb = stats(b.loc[a:bnd])
        print(f"  EW members {wname}: cagr {se['cagr']:6.1%} sh {se['sharpe']:5.2f} mdd {se['mdd']:6.1%}"
              f"   |  ETF blend: cagr {sb['cagr']:6.1%} sh {sb['sharpe']:5.2f}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--panel", default="both", choices=["old", "new", "both"])
    ap.add_argument("--top-n", type=int, default=20)
    ap.add_argument("--every", type=int, default=10)
    args = ap.parse_args()
    panels = ["old", "new"] if args.panel == "both" else [args.panel]
    out = []
    for panel in panels:
        print(f"\n=== {panel} panel ({args.top_n} names, every {args.every}) ===")
        p = load(panel)
        ew_universe(p)
        out.append(run(panel, args.top_n, args.every))
    df = pd.concat(out, ignore_index=True)
    os.makedirs(RESULTS, exist_ok=True)
    path = os.path.join(RESULTS, f"midsmall_scan_{args.every}.csv")
    df.to_csv(path, index=False)
    print(f"saved {path}")


if __name__ == "__main__":
    main()