"""Short-volume (FINRA REG SHO daily) factor scan on PIT US panels.

SVR = ShortVolume / TotalVolume per symbol-day. The REG SHO file for day D is
published after D's close; the scan therefore shifts all signals by 2 sessions
(conservative: usable at the close of D+2 at the earliest) before correlating
with forward returns.

Universes: research/lab panels us (S&P 500 ever-members) and midsmall
(S&P 400 + 600) — the mid/small cut is where a sub-$1M book has a capacity
edge over large funds. Output: monthly Spearman IC (1d/21d forward), quintile
spread means, and per-year ICs for the best factor.

    PYTHONPATH=. uv run python -m research.lab.short_volume [us|midsmall|all]
"""
from __future__ import annotations

import pickle
import sys

import numpy as np
import pandas as pd

SHORT_PATH = "data/research/short_volume.pkl"


def load_svr(panel_close: pd.DataFrame) -> pd.DataFrame:
    with open(SHORT_PATH, "rb") as fh:
        d = pickle.load(fh)
    short, total = d["short"], d["total"]
    cols = [c for c in panel_close.columns if c in short.columns]
    svr = (short[cols] / total[cols].replace(0.0, np.nan)).astype("float32")
    svr = svr.reindex(panel_close.index.union(svr.index)).sort_index()
    return svr


def factors(svr: pd.DataFrame) -> dict[str, pd.DataFrame]:
    f = {}
    f["svr5"] = svr.rolling(5, min_periods=3).mean()
    f["svr21"] = svr.rolling(21, min_periods=10).mean()
    m = svr.rolling(252, min_periods=60).mean()
    s = svr.rolling(252, min_periods=60).std()
    f["svr_z"] = (f["svr21"] - m) / s.replace(0.0, np.nan)
    f["svr_chg"] = f["svr5"] / f["svr21"].replace(0.0, np.nan) - 1.0
    return f


def scan(name: str, close: pd.DataFrame, member: pd.DataFrame, svr: pd.DataFrame) -> None:
    fac = {k: v.reindex(close.index) for k, v in factors(svr).items()}
    close = close.reindex(columns=svr.columns).ffill()
    member = member.reindex_like(close).fillna(False).astype(bool)
    fwd1 = close.shift(-1) / close - 1.0
    fwd21 = close.shift(-21) / close - 1.0
    lag = 2
    print(f"\n=== {name}: {close.shape[1]} symbols, "
          f"{close.index.min().date()}..{close.index.max().date()} ===")
    rows = []
    months = fwd1.index.to_period("M")
    for fname, fraw in fac.items():
        f = fraw.shift(lag).where(member)
        ics1, ics21 = [], []
        for per in sorted(set(months)):
            idx = months == per
            xdf = f.loc[idx]
            if xdf.notna().sum(axis=1).max() < 50:
                continue
            for ydf, out in ((fwd1.loc[idx], ics1), (fwd21.loc[idx], ics21)):
                mask = xdf.notna().to_numpy() & np.isfinite(ydf.to_numpy())
                if mask.sum() < 100:
                    continue
                xr = xdf.rank(axis=1).to_numpy()[mask]
                yr = ydf.where(mask).rank(axis=1).to_numpy()[mask]
                xm = xr - xr.mean()
                ym = yr - yr.mean()
                denom = np.sqrt((xm ** 2).sum() * (ym ** 2).sum())
                if denom > 0:
                    out.append(float((xm * ym).sum() / denom))
        ic1 = np.array(ics1)
        ic21 = np.array(ics21)
        rows.append({
            "factor": fname,
            "IC1_mean": ic1.mean(), "IC1_t": ic1.mean() / ic1.std() * np.sqrt(len(ic1)) if len(ic1) > 1 else np.nan,
            "IC21_mean": ic21.mean(), "IC21_t": ic21.mean() / ic21.std() * np.sqrt(len(ic21)) if len(ic21) > 1 else np.nan,
            "months": len(ic21)})
    print(pd.DataFrame(rows).round(4).to_string(index=False))

    # quintile spread, 21-day holding, gross
    for fname in ("svr21", "svr_chg", "svr_z"):
        f = fac[fname].shift(lag).where(member)
        pairs = []
        dates = f.index
        for i in range(0, len(dates) - 21, 21):
            t0 = dates[i]
            x = f.loc[t0].dropna()
            if len(x) < 50:
                continue
            q = pd.qcut(x.rank(method="first"), 5, labels=False)
            lo, hi = q[q == 0].index, q[q == 4].index
            r = close.iloc[i + 21] / close.iloc[i] - 1.0
            pairs.append((t0, float(r[hi].mean() - r[lo].mean())))
        s = np.array([p[1] for p in pairs])
        ys = pd.Series(s, index=pd.DatetimeIndex([p[0] for p in pairs]))
        yearly = (1 + ys).groupby(ys.index.year).prod() - 1
        print(f"{fname}: Q5-Q1 21d spread mean {s.mean()*100:.2f}%  "
              f"t={s.mean()/s.std()*np.sqrt(len(s)):.2f}  n={len(s)}  "
              f"yearly: " + " ".join(f"{y}:{v*100:+.0f}%" for y, v in yearly.items()))


def main() -> None:
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    if which in ("us", "all"):
        with open("data/research/us.panel.pkl", "rb") as fh:
            p = pickle.load(fh)
        scan("S&P 500 ever-members", p.close, p.member, load_svr(p.close))
    if which in ("midsmall", "all"):
        with open("data/research/us_midsmall.panel.pkl", "rb") as fh:
            p = pickle.load(fh)
        scan("S&P 400+600 (midsmall)", p.close, p.member, load_svr(p.close))


if __name__ == "__main__":
    main()