quant-trading/research/lab/tiingo_merge.py

"""Overlay Tiingo EOD caches (data/research/tiingo/) onto the US research panels.

S&P 500 -- ``apply_sp500(raw)``: fills the Yahoo-built panel's gaps, whole
missing columns for delisted names and late tails, keeping Yahoo cells where
both sources have data. Safety per ticker: >=50% of its PIT membership
sessions must be covered; where both sources overlap (>=20 sessions) the
median Tiingo/Yahoo close ratio must sit within 2% (same total-return
adjustment basis; a recycled ticker fails this). Two rename splices
(GPS<-GAP, DISCA<-WBD) backfill pre-rename eras from the successor ticker.
Read-only on the Yahoo artifact: deleting data/research/tiingo/ reverts it.
The merge prints an audit report incl. terminal gaps (last bar vs membership
end), which is the delisted-terminal check.

Mid/small -- ``write_yahoo_mid()``: converts the caches to the yahoo_mid CSV
format (split-adjusted OHLC with a Stock Splits column so load_bars can
recover the as-traded close, Adj Close as served) under
data/research/ext/tiingo_mid/, which fetch_us_mid.load_bars falls back to.

    PYTHONPATH=. uv run python -m research.lab.tiingo_merge --check
    PYTHONPATH=. uv run python -m research.lab.tiingo_merge --yahoo-mid
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np
import pandas as pd

DIR = os.path.join("data", "research", "tiingo")
MID_DIR = os.path.join("data", "research", "ext", "tiingo_mid")
FIELDS = ("open", "high", "low", "close", "volume")
SPLICE = {"GPS": "GAP", "DISCA": "WBD"}
MIN_COVERAGE = 0.5
MIN_OVERLAP = 20
RATIO_TOL = 0.02
TERMINAL_GAP_DAYS = 45
# Pure ticker changes where the successor's series is the same security's
# continuous price history (Yahoo/Tiingo migrate it on rename). NOT for
# merger-shell cases (RTN->RTX pre-2020 is United Technologies, not Raytheon).
SPLICE.update({"WLTW": "WTW", "ADS": "BFH", "CCE": "CCEP",
               "ESV": "VAL", "FBHS": "FBIN", "HFC": "DINO"})

_intervals_cache: dict | None = None


def _intervals() -> dict:
    global _intervals_cache
    if _intervals_cache is None:
        with open(os.path.join("data", "sp500_history.json")) as fh:
            _intervals_cache = json.load(fh)["intervals"]
    return _intervals_cache


def _cached(t: str) -> pd.DataFrame:
    h = pd.read_csv(os.path.join(DIR, f"{t.upper()}.csv"), index_col=0, parse_dates=True)
    return h[~h.index.duplicated(keep="last")].sort_index()


def cached_tickers() -> list[str]:
    if not os.path.isdir(DIR):
        return []
    return sorted(f[:-4] for f in os.listdir(DIR) if f.endswith(".csv"))


def _frame(t: str) -> pd.DataFrame:
    """Adjusted OHLCV in the panel convention (auto_adjust=True)."""
    h = _cached(t)
    out = pd.DataFrame({f: h["adj" + f[0].upper() + f[1:]] for f in FIELDS})
    return out


def _member_mask(t: str, dates: pd.DatetimeIndex) -> np.ndarray:
    m = np.zeros(len(dates), bool)
    for s, e in _intervals().get(t, []):
        si = dates.searchsorted(pd.Timestamp(s) if s else dates[0], side="left")
        ei = dates.searchsorted(pd.Timestamp(e) if e else dates[-1] + pd.Timedelta(days=1),
                                side="left")
        m[si:ei] = True
    return m


def _coverage(t: str, frame: pd.DataFrame, dates: pd.DatetimeIndex) -> float:
    mask = _member_mask(t, dates)
    if not mask.any():
        return 0.0
    have = frame["close"].reindex(dates).notna().to_numpy()
    return float((mask & have).sum()) / float(mask.sum())


def _cross_check(yahoo: pd.Series, tiingo: pd.Series) -> float | None:
    """Median Tiingo/Yahoo close ratio over the overlap; None if too small."""
    y = yahoo.reindex(tiingo.index)
    both = y.notna() & tiingo["close"].notna()
    if both.sum() < MIN_OVERLAP:
        return None
    return float((tiingo.loc[both, "close"] / y[both]).median())


def apply_sp500(raw: dict, *, verbose: bool = True) -> dict:
    """Return a copy of the raw panel dict with Tiingo cells filled in."""
    dates = raw["close"].index
    out = {k: v.copy() for k, v in raw.items()}
    merged, extended, replaced, skipped, rejected, spliced, terminals = [], [], [], [], [], [], []
    frames = {t: _frame(t) for t in cached_tickers()}

    def fill_column(t: str, frame: pd.DataFrame, mask: np.ndarray | None = None):
        added = t not in out["close"].columns
        for f in FIELDS:
            col = frame[f].reindex(dates)
            if mask is not None:
                col = col.where(mask)
            if added:
                out[f][t] = col
            else:
                out[f][t] = out[f][t].where(out[f][t].notna(), col)
        return added

    for t, frame in frames.items():
        cov = _coverage(t, frame, dates)
        if cov < MIN_COVERAGE:
            skipped.append((t, round(cov, 3)))
            continue
        replace = False
        if t in out["close"].columns:
            mask = _member_mask(t, dates)
            have = out["close"][t].reindex(dates).notna().to_numpy()
            if mask.any() and (mask & have).sum() / mask.sum() < MIN_COVERAGE:
                # a recycled ticker or a broken stub sits in the column: the
                # full series wins wherever it has data
                replace = True
            else:
                med = _cross_check(out["close"][t], frame)
                if med is not None and not (1 - RATIO_TOL < med < 1 + RATIO_TOL):
                    rejected.append((t, round(med, 4)))
                    continue
        if replace:
            for f in FIELDS:
                col = frame[f].reindex(dates)
                out[f][t] = col.where(col.notna(), out[f][t])
            replaced.append(t)
        else:
            added = fill_column(t, frame)
            (merged if added else extended).append(t)
        # terminal check: last bar vs membership end within the panel window
        col = out["close"][t].dropna()
        ends = [pd.Timestamp(e) if e else dates[-1] + pd.Timedelta(days=1)
                for _, e in _intervals().get(t, [])]
        mend = min(max(ends), dates[-1]) if ends else None
        if mend is not None and len(col):
            gap = (mend - col.index[-1]).days
            if gap > TERMINAL_GAP_DAYS:
                terminals.append((t, str(col.index[-1].date()), str(mend.date()), gap))

    for old, new in SPLICE.items():
        frame = frames.get(new)
        if frame is None and new in out["close"].columns:
            frame = pd.DataFrame({f: out[f][new] for f in FIELDS})
        if frame is None:
            continue
        mask = _member_mask(old, dates)
        if not mask.any():
            continue
        have = out["close"][old].reindex(dates).notna().to_numpy() if old in out["close"] else np.zeros(len(dates), bool)
        if (mask & have).sum() / mask.sum() >= MIN_COVERAGE:
            continue
        fill_column(old, frame, mask)
        spliced.append(f"{old}<-{new}")

    if verbose:
        print(f"tiingo overlay: {len(frames)} cached files; merged {len(merged)}, "
              f"extended {len(extended)}, replaced {len(replaced)}, spliced {len(spliced)}, "
              f"skipped {len(skipped)}, rejected {len(rejected)}")
        if merged:
            print(f"  merged: {', '.join(merged)}")
        if extended:
            print(f"  extended: {', '.join(extended)}")
        if replaced:
            print(f"  replaced (column was a stub/recycled): {', '.join(replaced)}")
        if spliced:
            print(f"  spliced: {', '.join(spliced)}")
        if skipped:
            print(f"  skipped (coverage<{MIN_COVERAGE}): "
                  + ", ".join(f"{t}:{c}" for t, c in skipped[:30])
                  + (" ..." if len(skipped) > 30 else ""))
        if rejected:
            print(f"  rejected (ratio off >{RATIO_TOL:.0%}): "
                  + ", ".join(f"{t}:{r}" for t, r in rejected))
        if terminals:
            print(f"  terminal gaps >{TERMINAL_GAP_DAYS}d: "
                  + ", ".join(f"{t} last={d} member_end={m} ({g}d)" for t, d, m, g in terminals))
    return out


PROD_FIELDS = ("open", "high", "low", "volume")


def _production_plan() -> tuple:
    """Which cached tickers to merge into the production store, and how."""
    import data_manager
    base = data_manager.load_raw("us")
    if base is None:
        raise FileNotFoundError("data/us.csv not found")
    dates = base.index
    frames = {t: _frame(t) for t in cached_tickers()}
    add, extend, replace, skip, rejected = [], [], [], [], []
    for t, fr in frames.items():
        cov = _coverage(t, fr, dates)
        if cov < MIN_COVERAGE:
            skip.append((t, round(cov, 3)))
            continue
        if t in base.columns:
            mask = _member_mask(t, dates)
            have = base[t].reindex(dates).notna().to_numpy()
            if mask.any() and (mask & have).sum() / mask.sum() < MIN_COVERAGE:
                replace.append(t)
                continue
            med = _cross_check(base[t], fr)
            if med is not None and not (1 - RATIO_TOL < med < 1 + RATIO_TOL):
                rejected.append((t, round(med, 4)))
                continue
            n = int(base[t].notna().sum())
            (extend if n else add).append(t)
        else:
            add.append(t)
    for old, new in SPLICE.items():
        if old in add:
            continue
        source = frames.get(new)
        if source is None and new in base.columns:
            source = pd.DataFrame({f: data_manager.load_raw("us", f)[new]
                                   for f in ("close",) + PROD_FIELDS})
        if source is None:
            continue
        mask = _member_mask(old, dates)
        if not mask.any():
            continue
        have = (base[old].reindex(dates).notna().to_numpy() if old in base.columns
                else np.zeros(len(dates), bool))
        if (mask & have).sum() / mask.sum() < MIN_COVERAGE:
            add.append(old)
            fr = source.reindex(dates)
            frames[old] = fr.where(np.broadcast_to(mask[:, None], fr.shape))
    return base, frames, add, extend, replace, skip, rejected


def apply_production(*, dry_run: bool = False) -> None:
    """Merge Tiingo frames into data/us*.csv (backup .bak-tiingo on first write)."""
    import data_manager
    from data_manager import _data_path, _panel_lock, _write_atomic
    base, frames, add, extend, replace, skip, rejected = _production_plan()
    names = add + extend + replace
    print(f"production merge: {len(frames)} cached; add {len(add)}, extend {len(extend)}, "
          f"replace {len(replace)}, skip {len(skip)}, rejected {len(rejected)}")
    if add:
        print(f"  add: {', '.join(add)}")
    if extend:
        print(f"  extend: {', '.join(extend)}")
    if replace:
        print(f"  replace: {', '.join(replace)}")
    if rejected:
        print(f"  rejected: {', '.join(f'{t}:{r}' for t, r in rejected)}")
    if dry_run or not names:
        return
    with _panel_lock("us"):
        fields = {f: data_manager.load_raw("us", f) for f in PROD_FIELDS}
        repl = set(replace)
        for f in FIELDS:
            panel = base if f == "close" else fields[f]
            if panel is None:
                continue
            for t in names:
                col = frames[t][f].reindex(panel.index)
                if t not in panel.columns:
                    panel[t] = col
                elif t in repl:
                    panel[t] = col.where(col.notna(), panel[t])
                else:
                    panel[t] = panel[t].where(panel[t].notna(), col)
            path = _data_path("us", f)
            if not os.path.exists(path + ".bak-tiingo"):
                import shutil
                shutil.copy2(path, path + ".bak-tiingo")
            _write_atomic(panel, path)
        print(f"production store updated for {len(names)} tickers "
              f"(backups: data/us*.csv.bak-tiingo)")


def write_yahoo_mid(tickers: list[str] | None = None) -> None:
    """Convert caches to yahoo_mid-format CSVs (split-adjusted OHLC + Stock
    Splits + as-traded close recoverable via load_bars)."""
    os.makedirs(MID_DIR, exist_ok=True)
    names = tickers if tickers is not None else cached_tickers()
    written = 0
    for t in names:
        h = _cached(t)
        if not {"close", "adjClose", "splitFactor"}.issubset(h.columns):
            print(f"{t}: missing columns; refetch with the current fetch_tiingo")
            continue
        splits = h["splitFactor"].replace(0, 1.0).fillna(1.0)
        later = splits[::-1].cumprod()[::-1].shift(-1).fillna(1.0)
        out = pd.DataFrame({
            "Open": h["open"] / later, "High": h["high"] / later,
            "Low": h["low"] / later, "Close": h["close"] / later,
            "Adj Close": h["adjClose"], "Volume": h["adjVolume"],
            "Dividends": h.get("divCash", pd.Series(0.0, index=h.index)).fillna(0.0),
            "Stock Splits": splits,
        })
        out.index.name = "Date"
        out.to_csv(os.path.join(MID_DIR, f"{t.upper()}.csv"))
        written += 1
    print(f"wrote {written} converted CSVs to {MID_DIR}")


def main() -> None:
    if "--yahoo-mid" in sys.argv:
        args = [a for a in sys.argv[1:] if not a.startswith("--")]
        write_yahoo_mid(args or None)
        return
    if "--production" in sys.argv:
        apply_production(dry_run="--apply" not in sys.argv)
        return
    import pickle
    with open(os.path.join("data", "research", "us_ohlcv.pkl"), "rb") as fh:
        raw = pickle.load(fh)
    apply_sp500(raw)


if __name__ == "__main__":
    main()