signal-platform/scripts/prep_survivorship_snapshot.py

"""Build the two survivorship-audit snapshots from a base snapshot.

Arm A is the universe the backtest already runs on. Arm B is the same universe
plus the S&P 500 constituents removed while the snapshot window was open, so the
strategy can trade the names that later left the index — including the ones that
failed.

The base file is never modified: both arms are copies, and the base hash is
re-checked at the end.

Three things this script has to get right, all of them learned from the data:

* **Schema drift.** The snapshot files on disk predate migrations 026 and 032,
  so ``select(Ticker)`` against them fails with ``no such column: tickers.cik``
  — and that is the path ``backtest_service._fetch_columns`` takes. Both copies
  are synced against ``Base.metadata`` before anything else touches them.
* **Frozen tails.** Alpaca does not end a halted symbol's series; it repeats the
  last price at zero volume. SBNY sits at 70.00 for 509 sessions after the FDIC
  closed it, and INFO sits at 108.61 for 1,090 sessions before a different
  company reuses the ticker. Only the symbols with a *reviewed cutoff* in the
  dataset are truncated. Everywhere else a long zero-volume run is a tripwire:
  the symbol is left out and reported for review rather than silently cut.
* **Terminal marks.** A failed bank stops trading mid-collapse, so the exchange
  bars never show the loss. A mark is written only where the dataset carries an
  explicit confirmation *and* the retained final bar matches the reviewed
  ``last_real_bar``. The mark is an assumption, not an observed execution.

Nothing is inferred. ``delisted_on`` / ``delisted_reason`` are written only for
confirmed terminal events; a series that merely ends early leaves them NULL and
is reported as an unresolved ending status.

Usage::

    python scripts/prep_survivorship_snapshot.py
    python scripts/prep_survivorship_snapshot.py --base backtest_snapshots/prod.sqlite --force
"""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import shutil
import sqlite3
import sys
from datetime import date, datetime, time, timedelta, timezone
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

# Bars are cut at the first run of this many consecutive zero-volume sessions.
ZERO_VOLUME_RUN = 5
# Close assigned to a confirmed terminal event. Not zero: a zero price divides
# through several return calculations, and 1 cent is a near-total loss either way.
TERMINAL_CLOSE = 0.01

DEFAULT_BASE = ROOT / "backtest_snapshots" / "prod.sqlite"
DEFAULT_REMOVALS = ROOT / "scripts" / "data" / "sp500_removals.json"


# ---------------------------------------------------------------------------
# schema sync
# ---------------------------------------------------------------------------

def sync_schema(path: Path) -> dict[str, list[str]]:
    """Add every column the models declare that the snapshot file lacks.

    Nullable and untyped-default only: this brings an old snapshot up to the
    current models without inventing data. Tables the snapshot does not have at
    all are left alone — a snapshot is a subset of production by design.
    """
    from sqlalchemy.dialects import sqlite as sqlite_dialect

    from app.database import Base
    import app.models  # noqa: F401  (registers every model on Base.metadata)

    dialect = sqlite_dialect.dialect()
    added: dict[str, list[str]] = {}
    con = sqlite3.connect(path)
    try:
        present = {r[0] for r in con.execute("SELECT name FROM sqlite_master WHERE type='table'")}
        for table in Base.metadata.sorted_tables:
            if table.name not in present:
                continue
            have = {r[1] for r in con.execute(f'PRAGMA table_info("{table.name}")')}
            for column in table.columns:
                if column.name in have:
                    continue
                type_sql = column.type.compile(dialect)
                con.execute(f'ALTER TABLE "{table.name}" ADD COLUMN "{column.name}" {type_sql}')
                added.setdefault(table.name, []).append(column.name)
        con.commit()
    finally:
        con.close()
    return added


# ---------------------------------------------------------------------------
# price data
# ---------------------------------------------------------------------------

def frozen_run_start(bars: list[Any]) -> int | None:
    """Index where the first run of >= ZERO_VOLUME_RUN zero-volume bars begins.

    A detector, not a cutter: the caller decides what to do with the answer, and
    only a reviewed cutoff is allowed to shorten a series.
    """
    volumes = [int(b.volume) for b in bars]
    i = 0
    while i < len(volumes):
        if volumes[i] == 0:
            j = i
            while j < len(volumes) and volumes[j] == 0:
                j += 1
            if j - i >= ZERO_VOLUME_RUN:
                return i
            i = j
        else:
            i += 1
    return None


def apply_reviewed_cutoff(bars: list[Any], cutoff: date) -> tuple[list[Any], int]:
    """Keep bars up to and including *cutoff*. Returns (kept, dropped)."""
    kept = [b for b in bars if b.timestamp.date() <= cutoff]
    return kept, len(bars) - len(kept)


async def fetch_bars(symbol: str, start: date, end: date) -> list[Any]:
    """Daily SPLIT-adjusted bars, matching how the snapshot's existing bars were ingested."""
    from app.config import settings
    from app.providers.alpaca import AlpacaOHLCVProvider

    if not settings.alpaca_api_key or not settings.alpaca_api_secret:
        raise SystemExit("Alpaca credentials are not configured; cannot restore price history.")

    from alpaca.data.enums import Adjustment
    from alpaca.data.requests import StockBarsRequest
    from alpaca.data.timeframe import TimeFrame

    provider = AlpacaOHLCVProvider(settings.alpaca_api_key, settings.alpaca_api_secret)
    alpaca_symbol = provider._to_alpaca_symbol(symbol)
    request = StockBarsRequest(
        symbol_or_symbols=alpaca_symbol,
        timeframe=TimeFrame.Day,
        start=datetime.combine(start, time.min, tzinfo=timezone.utc),
        end=datetime.combine(end + timedelta(days=1), time.min, tzinfo=timezone.utc),
        adjustment=Adjustment.SPLIT,
    )
    # 85 sequential requests are enough for Alpaca to drop one connection, and a
    # half-built arm is worse than a slow one.
    last_error: Exception | None = None
    for attempt in range(4):
        try:
            bars = await asyncio.to_thread(provider._client.get_stock_bars, request)
            return list(bars.data.get(alpaca_symbol, []))
        except Exception as exc:  # noqa: BLE001 - retry any transport failure
            last_error = exc
            if attempt < 3:
                await asyncio.sleep(2 * (attempt + 1))
    raise SystemExit(f"Alpaca fetch failed for {symbol} after 4 attempts: {last_error}")


# ---------------------------------------------------------------------------
# snapshot writes
# ---------------------------------------------------------------------------

def snapshot_window(path: Path) -> tuple[date, date]:
    con = sqlite3.connect(path)
    try:
        lo, hi = con.execute("SELECT min(date), max(date) FROM ohlcv_records").fetchone()
    finally:
        con.close()
    return date.fromisoformat(lo), date.fromisoformat(hi)


def trading_calendar(path: Path) -> list[date]:
    """SPY sessions from the snapshot's own benchmark series."""
    con = sqlite3.connect(path)
    try:
        rows = con.execute(
            "SELECT date FROM benchmark_prices WHERE symbol = 'SPY' ORDER BY date"
        ).fetchall()
    finally:
        con.close()
    return [date.fromisoformat(r[0]) for r in rows]


def next_session(calendar: list[date], after: date) -> date | None:
    for d in calendar:
        if d > after:
            return d
    return None


def insert_symbol(
    con: sqlite3.Connection,
    symbol: str,
    name: str,
    bars: list[Any],
    *,
    delisted_on: str | None,
    delisted_reason: str | None,
    terminal_date: date | None,
) -> int:
    """Insert one restored ticker and its bars. Returns the bar count written."""
    now = datetime.utcnow().isoformat(sep=" ")
    cur = con.execute(
        "INSERT INTO tickers (symbol, name, created_at, delisted_on, delisted_reason) "
        "VALUES (?, ?, ?, ?, ?)",
        (symbol, name, now, delisted_on, delisted_reason),
    )
    ticker_id = int(cur.lastrowid)

    rows = [
        (
            ticker_id,
            b.timestamp.date().isoformat(),
            float(b.open),
            float(b.high),
            float(b.low),
            float(b.close),
            int(b.volume),
            now,
        )
        for b in bars
    ]
    if terminal_date is not None:
        rows.append((
            ticker_id, terminal_date.isoformat(),
            TERMINAL_CLOSE, TERMINAL_CLOSE, TERMINAL_CLOSE, TERMINAL_CLOSE, 0, now,
        ))
    con.executemany(
        "INSERT INTO ohlcv_records "
        "(ticker_id, date, open, high, low, close, volume, created_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        rows,
    )
    return len(rows)


# ---------------------------------------------------------------------------
# driver
# ---------------------------------------------------------------------------

def file_hash(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


async def build(args: argparse.Namespace) -> int:
    base = Path(args.base).resolve()
    if not base.exists():
        raise SystemExit(f"Base snapshot not found: {base}")

    out_dir = Path(args.out_dir).resolve()
    arm_a = out_dir / args.arm_a_name
    arm_b = out_dir / args.arm_b_name
    events_path = arm_b.with_name(arm_b.stem + "-terminal-events.json")

    doc = json.loads(Path(args.removals).read_text(encoding="utf-8"))
    removals = doc["removals"]
    confirmations = {c["symbol"]: c for c in doc["terminal_confirmations"] if c.get("confirmed")}
    cutoffs = {
        c["symbol"]: date.fromisoformat(c["last_real_bar"])
        for c in doc.get("reviewed_cutoffs", {}).get("cutoffs", [])
    }
    # A corporate action ends a listing; a market-cap removal does not, so a
    # series that stops early without one is an ending we cannot account for.
    ENDING_EXPLAINED = {"acquisition", "receivership", "bankruptcy", "ticker_change"}

    base_hash_before = file_hash(base)
    window_start, window_end = snapshot_window(base)
    print(f"base     : {base}")
    print(f"window   : {window_start} .. {window_end}")
    print(f"removals : {len(removals)} rows, {sum(1 for r in removals if r['restore_symbol'])} restorable")

    for target in (arm_a, arm_b):
        if target.exists() and not args.force:
            raise SystemExit(f"{target} already exists (use --force to rebuild)")
        if target.exists():
            target.unlink()
        shutil.copy2(base, target)
        added = sync_schema(target)
        detail = ", ".join(f"{t}: {', '.join(c)}" for t, c in added.items()) or "nothing missing"
        print(f"copied   : {target.name}  (schema sync -> {detail})")

    calendar = trading_calendar(base)
    con = sqlite3.connect(arm_b)
    existing = {r[0].upper() for r in con.execute("SELECT symbol FROM tickers")}

    restored: list[dict] = []
    unresolved: list[dict] = []
    terminal_events: dict[str, str] = {}

    try:
        for row in removals:
            symbol = row["restore_symbol"]
            if not symbol:
                continue
            if symbol.upper() in existing:
                unresolved.append({"kind": "skipped", "symbol": symbol, "why": "already present in the base snapshot"})
                continue

            bars = await fetch_bars(symbol, window_start, window_end)
            if not bars:
                unresolved.append({"kind": "coverage", "symbol": symbol, "why": "Alpaca returned no bars in the window"})
                print(f"  {symbol:<6} NO DATA  (unresolved coverage)")
                continue

            cutoff = cutoffs.get(symbol) or cutoffs.get(row["symbol"])
            frozen_at = frozen_run_start(bars)
            if cutoff is not None:
                kept, dropped = apply_reviewed_cutoff(bars, cutoff)
                if not kept or kept[-1].timestamp.date() != cutoff:
                    unresolved.append({
                        "symbol": symbol,
                        "kind": "coverage",
                        "why": f"reviewed cutoff {cutoff} is not a bar in the downloaded series",
                    })
                    print(f"  {symbol:<6} CUTOFF MISSING {cutoff}  (unresolved coverage)")
                    continue
            elif frozen_at is not None:
                # Detector only. An unreviewed frozen run means the data changed
                # under us; restoring either half would be a guess, so report it.
                unresolved.append({
                    "kind": "coverage",
                    "symbol": symbol,
                    "why": (
                        f"unreviewed frozen run of >= {ZERO_VOLUME_RUN} zero-volume sessions "
                        f"from {bars[frozen_at].timestamp.date()}; needs a reviewed cutoff "
                        "before this symbol can be restored"
                    ),
                })
                print(f"  {symbol:<6} FROZEN RUN at {bars[frozen_at].timestamp.date()}  (unresolved, not restored)")
                continue
            else:
                kept, dropped = bars, 0

            effective_last = kept[-1].timestamp.date()
            confirmation = confirmations.get(row["symbol"]) or confirmations.get(symbol)
            terminal_date = None
            if confirmation is not None:
                confirmed_last = date.fromisoformat(confirmation["last_real_bar"])
                if effective_last != confirmed_last:
                    unresolved.append({
                        "symbol": symbol,
                        "kind": "terminal_mark",
                        "why": (
                            f"retained final bar {effective_last} does not match the reviewed "
                            f"last_real_bar {confirmed_last}; no terminal mark written"
                        ),
                    })
                    print(f"  {symbol:<6} MARK SKIPPED  final bar {effective_last} != confirmed {confirmed_last}")
                else:
                    terminal_date = next_session(calendar, effective_last)
                    if terminal_date is None or terminal_date > window_end:
                        terminal_date = None
                        unresolved.append({
                            "symbol": symbol,
                            "kind": "terminal_mark",
                            "why": "confirmed terminal event has no trading session left inside the window",
                        })
                    else:
                        terminal_events[symbol] = terminal_date.isoformat()

            # Documentary only, and only where the event is confirmed. A series
            # that merely ends early is an unresolved ending, not a delisting.
            delisted_on = None
            delisted_reason = None
            if terminal_date is not None:
                delisted_on = terminal_date.isoformat()
                delisted_reason = row["reason"]
            elif effective_last < window_end and row["reason"] not in ENDING_EXPLAINED:
                unresolved.append({
                    "kind": "ending_status",
                    "symbol": symbol,
                    "why": (
                        f"price history ends {effective_last}, before the window closes on "
                        f"{window_end}, and the source row ({row['reason']}) does not explain "
                        "an ending; ending status unresolved and delisting fields left null"
                    ),
                })

            written = insert_symbol(
                con, symbol, row["name"], kept,
                delisted_on=delisted_on, delisted_reason=delisted_reason,
                terminal_date=terminal_date,
            )
            restored.append({
                "symbol": symbol,
                "removed_on": row["removed_on"],
                "reason": row["reason"],
                "bars": written,
                "first_bar": kept[0].timestamp.date().isoformat(),
                "effective_last_bar": effective_last.isoformat(),
                "frozen_bars_dropped": dropped,
                "terminal_mark": terminal_events.get(symbol),
            })
            mark = f"  terminal={terminal_events[symbol]} @ {TERMINAL_CLOSE}" if symbol in terminal_events else ""
            drop = f"  dropped={dropped}" if dropped else ""
            print(f"  {symbol:<6} bars={written:<5} {kept[0].timestamp.date()}..{effective_last}{drop}{mark}")

        con.commit()
    finally:
        con.close()

    events_path.write_text(json.dumps(terminal_events, indent=2) + "\n", encoding="utf-8")

    if file_hash(base) != base_hash_before:
        raise SystemExit(f"BASE SNAPSHOT WAS MODIFIED: {base}")

    con_a = sqlite3.connect(arm_a)
    con_b = sqlite3.connect(arm_b)
    counts = {
        "arm_a_tickers": con_a.execute("SELECT count(*) FROM tickers").fetchone()[0],
        "arm_a_bars": con_a.execute("SELECT count(*) FROM ohlcv_records").fetchone()[0],
        "arm_b_tickers": con_b.execute("SELECT count(*) FROM tickers").fetchone()[0],
        "arm_b_bars": con_b.execute("SELECT count(*) FROM ohlcv_records").fetchone()[0],
    }
    con_a.close()
    con_b.close()

    by_kind: dict[str, list[dict]] = {}
    for row in unresolved:
        by_kind.setdefault(row.get("kind", "other"), []).append(row)

    print("\nsummary")
    print(f"  restored           : {len(restored)}")
    print(f"  unresolved coverage: {len(by_kind.get('coverage', []))}  (symbols not restored)")
    print(f"  unresolved endings : {len(by_kind.get('ending_status', []))}  (restored, ending status unknown)")
    print(f"  terminal marks     : {len(terminal_events)} -> {events_path.name}")
    print(f"  arm A              : {counts['arm_a_tickers']} tickers, {counts['arm_a_bars']} bars")
    print(f"  arm B              : {counts['arm_b_tickers']} tickers, {counts['arm_b_bars']} bars")
    print(f"  base unchanged     : sha256 {base_hash_before[:16]}...")
    for kind in ("coverage", "ending_status", "terminal_mark"):
        for row in by_kind.get(kind, []):
            print(f"  UNRESOLVED[{kind}] {row['symbol']:<6} {row['why']}")

    manifest = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "base": str(base),
        "base_sha256": base_hash_before,
        "window": {"start": window_start.isoformat(), "end": window_end.isoformat()},
        "zero_volume_run": ZERO_VOLUME_RUN,
        "terminal_close": TERMINAL_CLOSE,
        "counts": counts,
        "restored": restored,
        "unresolved": unresolved,
        "terminal_events": terminal_events,
    }
    manifest_path = arm_b.with_name(arm_b.stem + "-prep-manifest.json")
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(f"  manifest           : {manifest_path.name}")
    return 0


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--base", default=str(DEFAULT_BASE), help="Base snapshot to copy (never modified).")
    p.add_argument("--out-dir", default=str(DEFAULT_BASE.parent), help="Where the arm copies are written.")
    p.add_argument("--arm-a-name", default="survivorship-a.sqlite")
    p.add_argument("--arm-b-name", default="survivorship-b.sqlite")
    p.add_argument("--removals", default=str(DEFAULT_REMOVALS))
    p.add_argument("--force", action="store_true", help="Rebuild the arm files if they already exist.")
    return p.parse_args()


if __name__ == "__main__":
    raise SystemExit(asyncio.run(build(parse_args())))