"""Fails-to-deliver factor scan on PIT US panels.

FTD value / 20-day dollar volume = fails intensity: the share of a stock's
daily liquidity that is stuck unsettled. High intensity marks short-sale
constraint / squeeze fuel and, on the classic evidence, predicts negative
returns in small caps.

PIT: SEC publishes the CNMS fails files roughly 3 weeks after settlement;
the scan shifts the FTD panel forward by 21 calendar days before
reindexing to trading days (ffill, staleness-capped).

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

import pickle
import sys

import numpy as np
import pandas as pd

FAILS_PATH = "data/research/fails.pkl"
PUBLISH_LAG_DAYS = 21


def load_ftd(close: pd.DataFrame, volume: pd.DataFrame | None) -> dict[str, pd.DataFrame]:
    with open(FAILS_PATH, "rb") as fh:
        d = pickle.load(fh)
    shares, value = d["shares"], d["value"]
    cols = [c for c in close.columns if c in shares.columns]
    f = {}
    v = value[cols]
    f["ftd_value"] = v
    f["ftd_shares"] = shares[cols]
    if volume is not None and len(volume):
        vol = volume.reindex(columns=cols).ffill()
        px = close.reindex(columns=cols).ffill()
        adv = (px * vol).rolling(20, min_periods=10).mean()
        f["ftd_int"] = v / adv.replace(0.0, np.nan)
        sh = shares[cols]
        f["ftd_sh"] = sh / vol.rolling(20, min_periods=10).mean().replace(0.0, np.nan)
    return f


def scan(name: str, close: pd.DataFrame, member: pd.DataFrame,
         volume: pd.DataFrame | None) -> None:
    raw = load_ftd(close, volume)
    member = member.reindex_like(close).fillna(False)
    print(f"\n=== {name}: {close.shape[1]} symbols, "
          f"{close.index.min().date()}..{close.index.max().date()} ===")
    fwd1 = close.shift(-1) / close - 1.0
    fwd21 = close.shift(-21) / close - 1.0
    months = fwd1.index.to_period("M")
    rows = []
    lagged: dict[str, pd.DataFrame] = {}
    for fname, fraw in raw.items():
        f = fraw.copy()
        f.index = f.index + pd.Timedelta(days=PUBLISH_LAG_DAYS)
        f = (f.reindex(index=close.index.union(f.index), columns=close.columns)
             .ffill(limit=45).reindex(close.index))
        f = f.where(member)
        lagged[fname] = f
        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, ic21 = np.array(ics1), np.array(ics21)
        rows.append({"factor": fname,
                     "IC1": ic1.mean(), "t1": ic1.mean() / ic1.std() * np.sqrt(len(ic1)) if len(ic1) > 1 else np.nan,
                     "IC21": ic21.mean(), "t21": 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))

    for fname in [k for k in ("ftd_int", "ftd_sh") if k in lagged]:
        f = lagged[fname]
        pairs = []
        dates = f.index
        for i in range(0, len(dates) - 21, 10):
            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())))
        if not pairs:
            continue
        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)}  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)
        vol = None
        try:
            vol = pd.read_csv("data/us_volume.csv", index_col=0, parse_dates=True)
        except Exception:  # noqa: BLE001
            pass
        scan("S&P 500 ever-members", p.close, p.member, vol)
    if which in ("midsmall", "all"):
        with open("data/research/us_midsmall.panel.pkl", "rb") as fh:
            p = pickle.load(fh)
        vol = getattr(p, "volume", None)
        scan("S&P 400+600 (midsmall)", p.close, p.member, vol)


if __name__ == "__main__":
    main()