"""
Which strategies and factors carry alpha against QQQ, before and after 2017?

    PYTHONPATH=. uv run python -m research.lab.fetch_long_history       # once (QQQ)
    PYTHONPATH=. uv run --with numba python -m research.lab.qqq_scan [--part sleeves|factors|mega]

sleeves  the current sleeves (us_final, short_horizon) and TrendRider
factors  every price factor (ranks_us) and every EDGAR fundamental, long-only
         top 20 among PIT S&P 500 members, every 21 days, buffer 2
mega     the same idea inside a QQQ-like universe: the 100 members with the
         highest 126-day dollar volume (market caps are missing for many
         delisted names, dollar volume is not), top 10 / 20

Alpha is the intercept of daily excess returns on QQQ's excess return, with
a Newey-West t-stat (research/lab/alpha.py). Periods: design 2006-09..2016,
holdout 2017..2026-09. A factor whose alpha is positive in both, with t > 2
over the whole span, is a candidate; everything else is noise or beta.
Writes data/research/results/qqq_scan_{part}.pkl (daily returns) and .csv.
"""

from __future__ import annotations

import argparse
import dataclasses
import os
import pickle
import warnings

import pandas as pd

from research.lab import fundamentals as fu
from research.lab.alpha import alpha_stats, bench_returns, fmt_alpha
from research.lab.factor_scan import RESULTS, load
from research.lab.factors import _dollar_volume
from research.lab.rotation import factor_ranks
from research.lab.us_sweep import _blend, run, sector_rank

warnings.filterwarnings("ignore")
LONG = ("2006-09-27", "2026-09-24")
DESIGN = ("2006-09-27", "2016-12-31")
HOLDOUT = ("2017-01-01", "2026-09-24")


def report(name: str, r: pd.Series, rows: list, bench: pd.Series, spy: pd.Series) -> None:
    r = r.dropna()
    full = alpha_stats(r, bench.reindex(r.index))
    d = alpha_stats(r.loc[DESIGN[0]:DESIGN[1]], bench)
    h = alpha_stats(r.loc[HOLDOUT[0]:HOLDOUT[1]], bench)
    s = alpha_stats(r, spy.reindex(r.index))
    rows.append({"name": name, **{f"full_{k}": v for k, v in full.items()},
                 **{f"design_{k}": v for k, v in d.items()}, **{f"holdout_{k}": v for k, v in h.items()},
                 "spy_alpha": s["alpha"], "spy_t": s["t"]})
    print(f"{fmt_alpha(name, full)} | design a {d['alpha']*100:5.1f}% t {d['t']:5.2f} | "
          f"holdout a {h['alpha']*100:5.1f}% t {h['t']:5.2f} cagr {h['cagr']*100:5.1f}% vs {h['bench_cagr']*100:5.1f}% "
          f"| vs SPY a {s['alpha']*100:5.1f}% t {s['t']:5.2f}", flush=True)


def part_sleeves(bench, spy, rows) -> dict:
    R = pd.read_pickle(os.path.join(RESULTS, "us_final_sleeves.pkl"))
    extra = os.path.join(RESULTS, "short_horizon_sleeves.pkl")
    if os.path.exists(extra):
        R = R.join(pd.read_pickle(extra))
    with open(os.path.join("data", "research", "letf_history.pkl"), "rb") as fh:
        tr = pickle.load(fh)["equity"]
    R["trend_rider_v7_vt24 (synthetic)"] = tr["trend_rider_v7_vt24 [lag]"].pct_change().reindex(R.index)
    for c in R.columns:
        report(c, R[c], rows, bench, spy)
    return {}


def part_factors(p, bench, spy, rows) -> dict:
    pr = factor_ranks(p, "us")
    F = fu.load(p)
    scores = {k: v for k, v in pr.items()}
    for k, v in F.items():
        if k == "mcap":
            continue
        scores[f"f:{k}"] = v.where(p.member).rank(axis=1, pct=True)
        scores[f"f:{k} (sector)"] = sector_rank(v, p)
    out = {}
    for name, sc in scores.items():
        st = run(p, sc.where(p.member), [LONG], top_n=20, buffer=2.0, every=21)
        out[name] = st["_rets"]
        report(name, st["_rets"], rows, bench, spy)
    return out


def mega_panel(p, n: int = 100):
    dv = _dollar_volume(p).rolling(126, min_periods=63).mean().where(p.member)
    top = dv.rank(axis=1, ascending=False) <= n
    return dataclasses.replace(p, member=p.member & top)


def part_mega(p, bench, spy, rows) -> dict:
    pm = mega_panel(p)
    pr = factor_ranks(p, "us")
    F = fu.load(p)
    ew = pm.close.notna().astype(float).where(pm.member)
    scores = {
        "mega EW (100)": (ew, 10000),
        "mega overnight+rec": (_blend(pr["overnight_mom"], pr["rec_126"]), None),
        "mega mom_12_1": (pr["mom_12_1"], None), "mega resid_mom": (pr["resid_mom"], None),
        "mega high_52w": (pr["high_52w"], None), "mega sharpe_mom": (pr["sharpe_mom"], None),
        "mega overnight_mom": (pr["overnight_mom"], None), "mega rec_126": (pr["rec_126"], None),
        "mega mom_6_1": (pr["mom_6_1"], None), "mega trend_hzz": (pr["trend_hzz"], None),
        "mega low_beta": (pr["low_beta"], None), "mega ivol": (pr["ivol"], None),
        "mega gp_a": (F["gp_a"].rank(axis=1, pct=True), None), "mega roe": (F["roe"].rank(axis=1, pct=True), None),
        "mega mom_12_1+gp_a": (_blend(pr["mom_12_1"], F["gp_a"].rank(axis=1, pct=True)), None),
    }
    out = {}
    for name, (sc, n_all) in scores.items():
        for n in ([n_all] if n_all else (10, 20)):
            st = run(pm, sc.where(pm.member), [LONG], top_n=n, buffer=2.0, every=21)
            key = name if n_all else f"{name} top{n}"
            out[key] = st["_rets"]
            report(key, st["_rets"], rows, bench, spy)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--part", nargs="*", default=["sleeves", "mega", "factors"])
    args = ap.parse_args()
    p, _ = load("us")
    bench = bench_returns(p.dates)
    spy = bench_returns(p.dates, "SPY")
    report("QQQ itself", bench.loc[LONG[0]:], [], bench, spy)
    for part in args.part:
        print(f"\n=== {part}: alpha vs QQQ (full 2006-2026) | design 2006-16 | holdout 2017-26 | vs SPY", flush=True)
        rows = []
        if part == "sleeves":
            out = part_sleeves(bench, spy, rows)
        elif part == "factors":
            out = part_factors(p, bench, spy, rows)
        else:
            out = part_mega(p, bench, spy, rows)
        pd.DataFrame(rows).to_csv(os.path.join(RESULTS, f"qqq_scan_{part}.csv"), index=False)
        if out:
            pd.DataFrame(out).to_pickle(os.path.join(RESULTS, f"qqq_scan_{part}.pkl"))


if __name__ == "__main__":
    main()