"""Quantify price gaps against PIT membership in the US research panels and
emit a Tiingo backfill plan (free tier: 500 unique symbols/month, 50/hour).

For S&P 500 ever-members: names with no data, a late start or an early end
relative to their membership intervals, in both the production books
(data/us.csv) and the research panel (data/research/us_ohlcv.pkl). For the
mid/small panels: the Yahoo-missing tickers (data/research/ext/yahoo_mid/
missing.json) ranked by snapshot weight and recency, since only the top of
that list fits the monthly quota.

Writes data/research/tiingo_plan.json.

    PYTHONPATH=. uv run python -m research.lab.panel_gaps
"""
from __future__ import annotations

import json
import os
import pickle

import numpy as np
import pandas as pd

TOL_DAYS = 45
PLAN = os.path.join("data", "research", "tiingo_plan.json")


def sp500_intervals() -> dict:
    with open(os.path.join("data", "sp500_history.json")) as fh:
        return json.load(fh)["intervals"]


def classify(px: pd.DataFrame, intervals: dict) -> pd.DataFrame:
    """Per ever-member: membership vs data presence within the panel window."""
    rows = []
    pstart, pend = px.index[0], px.index[-1]
    for t, iv in intervals.items():
        iv = [(pd.Timestamp(s) if s else None, pd.Timestamp(e) if e else None) for s, e in iv]
        mstart = max(min((s for s, _ in iv if s), default=pstart), pstart)
        mend = min(max((e for _, e in iv if e), default=pend), pend)
        if mend <= pstart or mstart >= pend:
            continue
        col = px[t].dropna() if t in px.columns else pd.Series(dtype=float)
        first = col.index[0] if len(col) else None
        last = col.index[-1] if len(col) else None
        n_mem = int(((px.index >= mstart) & (px.index < mend)).sum())
        n_px = int(((col.index >= mstart) & (col.index < mend)).sum()) if len(col) else 0
        flags = []
        if not len(col):
            flags.append("no_data")
        else:
            if first > mstart + pd.Timedelta(days=TOL_DAYS):
                flags.append("late_start")
            if last < mend - pd.Timedelta(days=TOL_DAYS):
                flags.append("early_end")
        rows.append(dict(ticker=t, m_first=mstart.date(), m_last=mend.date(),
                         px_first=first.date() if first is not None else None,
                         px_last=last.date() if last is not None else None,
                         n_mem=n_mem, n_px=n_px,
                         cov=round(n_px / n_mem, 3) if n_mem else np.nan,
                         flags=",".join(flags)))
    return pd.DataFrame(rows)


def mid_small_plan() -> pd.DataFrame:
    """Rank the Yahoo-missing mid/small tickers: total snapshot weight, with
    recent years weighted fully and older ones decayed."""
    from research.lab.fetch_us_mid import load_holdings

    miss_path = os.path.join("data", "research", "ext", "yahoo_mid", "missing.json")
    missing = set(json.load(open(miss_path)))
    h = load_holdings()
    h = h[h["ticker"].isin(missing)]
    if h.empty:
        return h
    recency = pd.Series(1.0, index=h.index)
    recency[h["date"] < "2014-01-01"] = 0.25
    recency[(h["date"] >= "2014-01-01") & (h["date"] < "2019-01-01")] = 0.5
    h = h.assign(score=h["weight"] * recency)
    g = h.groupby("ticker")
    out = pd.DataFrame({
        "n_snaps": g.size(),
        "w_sum": g["weight"].sum().round(4),
        "score": g["score"].sum().round(4),
        "first": g["date"].min().dt.date,
        "last": g["date"].max().dt.date,
        "name": g["name"].first(),
    }).sort_values("score", ascending=False)
    return out


def main() -> None:
    intervals = sp500_intervals()

    import data_manager
    prod = data_manager.load_raw("us")
    print(f"production data/us.csv: {prod.shape[1]} cols {prod.index[0].date()}..{prod.index[-1].date()}")

    from research.lab.tiingo_merge import apply_sp500
    with open(os.path.join("data", "research", "us_ohlcv.pkl"), "rb") as fh:
        raw = pickle.load(fh)
    res = apply_sp500(raw, verbose=False)["close"]
    print(f"research us_ohlcv.pkl + tiingo: {res.shape[1]} cols "
          f"{res.index[0].date()}..{res.index[-1].date()}")

    plan: dict[str, list[str]] = {}
    for label, px in (("production", prod), ("research", res)):
        df = classify(px, intervals)
        bad = df[df["flags"] != ""].sort_values(["n_mem"], ascending=False)
        print(f"\n[{label}] ever-members with intervals in window: {len(df)}; with flags: {len(bad)}")
        for fl in ("no_data", "late_start", "early_end"):
            sub = bad[bad["flags"].str.contains(fl)]
            print(f"  {fl}: {len(sub)}")
            cols = ["ticker", "m_first", "m_last", "px_first", "px_last", "cov"]
            if len(sub):
                print(sub[cols].head(25).to_string(index=False))
        names = sorted(set(bad[bad["flags"].str.contains("no_data|early_end")]["ticker"]))
        plan[f"sp500_{label}"] = names
        print(f"  -> {len(names)} no_data/early_end names")

    ms = mid_small_plan()
    print(f"\n[mid/small] missing tickers with snapshot weight: {len(ms)}")
    print(ms.head(30).to_string())
    print(f"  weight coverage reachable in quota (top 500): "
          f"{ms['w_sum'].head(500).sum() / ms['w_sum'].sum():.1%}")
    plan["mid_small"] = list(ms.index)

    with open(PLAN, "w") as fh:
        json.dump(plan, fh, indent=1)
    print(f"\nSaved {PLAN}")


if __name__ == "__main__":
    main()