fithia2/apps/intraday_bt/scripts/build_daily_snapshot.py

"""Build a clean ORB daily-bar snapshot from Oracle daily bars.

This intentionally avoids intraday-cache fallback so sparse local intraday
fragments cannot be written as if they were full daily history.

Example:
    python -m apps.intraday_bt.scripts.build_daily_snapshot \
        --snapshot-id orb_daily_v49_20260508 \
        --universe broad \
        --start-date 2024-02-03 \
        --end-date 2026-05-08 \
        --force
"""
from __future__ import annotations

import argparse
import asyncio
import json
import shutil
import sys
import time
from pathlib import Path

import pyarrow.parquet as pq
import yaml

from apps.intraday_bt.oracle import make_intraday_oracle_client
from libs.common.config import get_settings
from libs.intraday.cache import DailyBarCache
from libs.intraday.screener import fetch_daily_bars_bulk


_UNIVERSE_FILES = {
    "broad": "configs/symbols_broad_snapshot_3408.yaml",
    "midlarge": "configs/symbols_midlarge_snapshot_exact.yaml",
    "largecap": "configs/symbols.yaml",
}

_DEFAULT_EXTRA_TICKERS = ("QQQ", "SPY", "TQQQ", "QQQM", "SGOV")


def _load_tickers(universe: str, extra_tickers: list[str]) -> list[str]:
    if universe in _UNIVERSE_FILES:
        path = Path(_UNIVERSE_FILES[universe])
        with path.open() as f:
            raw = yaml.safe_load(f) or []
        if isinstance(raw, dict):
            tickers = raw.get("tickers") or raw.get("symbols") or []
        else:
            tickers = raw
    else:
        path = Path(universe)
        with path.open() as f:
            raw = yaml.safe_load(f) or []
        tickers = raw.get("tickers") if isinstance(raw, dict) else raw

    seen: set[str] = set()
    result: list[str] = []
    for ticker in [*tickers, *extra_tickers]:
        symbol = str(ticker).strip().upper()
        if not symbol or symbol in seen:
            continue
        seen.add(symbol)
        result.append(symbol)
    return result


def _snapshot_summary(snapshot_dir: Path, start_date: str, end_date: str) -> dict[str, object]:
    files = sorted(snapshot_dir.glob("*.parquet"))
    rows_total = 0
    full_end = 0
    late_start = 0
    examples_late: list[dict[str, object]] = []
    examples_short: list[dict[str, object]] = []
    for path in files:
        try:
            table = pq.read_table(str(path))
            rows = table.to_pylist()
        except Exception:
            continue
        if not rows:
            continue
        rows_total += len(rows)
        first = str(rows[0]["date"])[:10]
        last = str(rows[-1]["date"])[:10]
        if last >= end_date:
            full_end += 1
        if first > start_date:
            late_start += 1
            if len(examples_late) < 10:
                examples_late.append(
                    {"ticker": path.stem, "rows": len(rows), "first": first, "last": last}
                )
        if len(rows) < 20 and len(examples_short) < 10:
            examples_short.append(
                {"ticker": path.stem, "rows": len(rows), "first": first, "last": last}
            )

    return {
        "files": len(files),
        "rows_total": rows_total,
        "full_end_files": full_end,
        "late_start_files": late_start,
        "examples_late_start": examples_late,
        "examples_short": examples_short,
    }


async def _run(args: argparse.Namespace) -> int:
    snapshot_dir = Path(args.output_root) / args.snapshot_id
    if snapshot_dir.exists():
        if not args.force:
            raise SystemExit(f"{snapshot_dir} already exists; use --force to replace it")
        shutil.rmtree(snapshot_dir)
    snapshot_dir.mkdir(parents=True, exist_ok=True)

    tickers = _load_tickers(args.universe, args.extra_ticker)
    if args.limit is not None:
        tickers = tickers[: args.limit]

    print(
        f"Building daily snapshot {args.snapshot_id}: {len(tickers)} tickers, "
        f"{args.start_date}..{args.end_date}, concurrency={args.concurrency}"
    )

    started = time.time()
    last_pct = [-1]

    def progress(completed: int, total: int) -> None:
        pct = int(completed / total * 100) if total else 100
        if pct >= last_pct[0] + 2 or completed == total:
            last_pct[0] = pct
            sys.stdout.write(f"\r  {completed}/{total} ({pct}%)")
            sys.stdout.flush()

    async with make_intraday_oracle_client(get_settings()) as client:
        bars = await fetch_daily_bars_bulk(
            tickers,
            args.start_date,
            args.end_date,
            client,
            cache=DailyBarCache(str(snapshot_dir)),
            intraday_cache_fallback=None,
            prefer_intraday_fallback=False,
            skip_oracle_when_unhealthy=False,
            concurrency=args.concurrency,
            progress_callback=progress,
        )
    print()

    summary = _snapshot_summary(snapshot_dir, args.start_date, args.end_date)
    summary.update(
        {
            "snapshot_id": args.snapshot_id,
            "universe": args.universe,
            "requested_tickers": len(tickers),
            "returned_tickers": len(bars),
            "missing_tickers": len(tickers) - len(bars),
            "start_date": args.start_date,
            "end_date": args.end_date,
            "elapsed_sec": round(time.time() - started, 2),
        }
    )
    manifest_path = snapshot_dir / "_manifest.json"
    manifest_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--snapshot-id", required=True)
    parser.add_argument("--universe", default="broad")
    parser.add_argument("--start-date", required=True)
    parser.add_argument("--end-date", required=True)
    parser.add_argument("--output-root", default="data/cache/daily_snapshots")
    parser.add_argument("--extra-ticker", action="append", default=list(_DEFAULT_EXTRA_TICKERS))
    parser.add_argument("--concurrency", type=int, default=20)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    return asyncio.run(_run(args))


if __name__ == "__main__":
    raise SystemExit(main())