fithia2/scripts/enrich_tier2_features.py

"""Enrich snapshot with Tier 2 features: FINRA short ratio, Hurst, Entropy, Sector momentum.

Usage:
    PYTHONUNBUFFERED=1 uv run python3 scripts/enrich_tier2_features.py \
        --input data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit \
        --output data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_tier2
"""
from __future__ import annotations

import argparse
import asyncio
import concurrent.futures
import datetime as dt
import json
import sys
from collections import defaultdict
from pathlib import Path

import asyncpg
import pyarrow as pa
import pyarrow.parquet as pq
import requests

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from libs.features.market_features import pre_event_entropy, pre_event_hurst
from libs.common.config import get_settings
from libs.oracle_client.models import PriceBar

ORACLE_URL = "http://localhost:18001"
FEATURES = ["pre_event_hurst_60d", "pre_event_entropy_60d", "pre_event_short_ratio", "pre_event_sector_momentum_20d"]
_PRICE_HISTORY_START = "2019-01-01"
_SHORT_RATIO_HISTORY_DAYS = 2000
_SHORT_RATIO_HISTORY_LOOKBACK_BUFFER_DAYS = 10
_FINRA_DAILY_MAX_CALENDAR_LOOKBACK_DAYS = 14
_FINRA_DAILY_URL_TEMPLATE = "https://cdn.finra.org/equity/regsho/daily/CNMSshvol{yyyymmdd}.txt"

SECTOR_ETF = {
    "Technology": "XLK", "Health Care": "XLV", "Financials": "XLF",
    "Consumer Discretionary": "XLY", "Industrials": "XLI",
    "Energy": "XLE", "Utilities": "XLU", "Real Estate": "XLRE",
    "Materials": "XLB", "Communication Services": "XLC",
    "Consumer Staples": "XLP", "Healthcare": "XLV",
}

SECTOR_ALIASES = {
    "Financial Services": "Financials",
    "Consumer Cyclical": "Consumer Discretionary",
    "Consumer Defensive": "Consumer Staples",
    "Basic Materials": "Materials",
}

# Cache sector ETF / index bars and FINRA history to keep enrich runs fast.
_sector_etf_cache: dict[str, list[dict]] = {}
_short_ratio_cache: dict[str, tuple[int, list[dict]]] = {}
_short_ratio_db_cache: dict[str, list[dict]] = {}
_finra_daily_cache: dict[str, dict[str, float]] = {}
_oracle_short_ratio_unavailable = False
_http = requests.Session()


def _fetch_bars_raw(ticker: str, start: str, end: str) -> list[dict]:
    resp = requests.get(f"{ORACLE_URL}/api/v1/price/data/{ticker}",
                       params={"start_date": start, "end_date": end}, timeout=30)
    if resp.status_code != 200:
        return []
    data = resp.json()
    return data.get("bars") or data.get("data") or []


def _fetch_full_bars_raw(ticker: str) -> list[dict]:
    if ticker not in _sector_etf_cache:
        _sector_etf_cache[ticker] = _fetch_bars_raw(
            ticker,
            _PRICE_HISTORY_START,
            dt.date.today().isoformat(),
        )
    return _sector_etf_cache[ticker]


def normalize_sector_name(sector: str | None) -> str | None:
    if sector is None:
        return None
    normalized = str(sector).strip()
    if not normalized:
        return None
    return SECTOR_ALIASES.get(normalized, normalized)


def _extract_short_ratio_history(payload: dict) -> list[dict]:
    history = payload.get("history")
    if isinstance(history, list):
        return history
    data = payload.get("data")
    if isinstance(data, list):
        return data
    return []


def _required_short_ratio_days(event_date: str) -> int:
    event_dt = dt.date.fromisoformat(event_date)
    raw_days = (dt.date.today() - event_dt).days + _SHORT_RATIO_HISTORY_LOOKBACK_BUFFER_DAYS
    return max(60, min(raw_days, _SHORT_RATIO_HISTORY_DAYS))


async def _query_short_ratio_history_from_db(ticker: str) -> list[dict]:
    dsn = get_settings().postgres_dsn.replace("+asyncpg", "")
    conn = await asyncpg.connect(dsn=dsn)
    try:
        rows = await conn.fetch(
        """
        select trade_date::text as date,
               short_volume::double precision / nullif(total_volume::double precision, 0.0) as short_ratio
        from short_sale_daily
        where ticker_raw = $1
          and total_volume is not null
          and total_volume > 0
        order by trade_date desc
        """,
            ticker,
        )
    finally:
        await conn.close()
    return [{"date": row["date"], "short_ratio": row["short_ratio"]} for row in rows if row["short_ratio"] is not None]


def _fetch_short_ratio_history_from_db(ticker: str) -> list[dict]:
    if ticker in _short_ratio_db_cache:
        return _short_ratio_db_cache[ticker]
    try:
        history = asyncio.run(_query_short_ratio_history_from_db(ticker))
    except Exception:
        history = []
    _short_ratio_db_cache[ticker] = history
    return history


def _fetch_short_ratio_history_from_oracle(ticker: str, days: int) -> list[dict]:
    global _oracle_short_ratio_unavailable
    if _oracle_short_ratio_unavailable:
        return []
    cached = _short_ratio_cache.get(ticker)
    if cached is not None:
        cached_days, cached_history = cached
        if cached_days >= days:
            return cached_history
    try:
        resp = _http.get(
            f"{ORACLE_URL}/api/v1/finra/short-ratio/{ticker}",
            params={"days": days},
            timeout=10,
        )
    except requests.RequestException:
        _oracle_short_ratio_unavailable = True
        _short_ratio_cache[ticker] = (days, [])
        return []
    if resp.status_code != 200:
        _short_ratio_cache[ticker] = (days, [])
        return []
    history = _extract_short_ratio_history(resp.json())
    _short_ratio_cache[ticker] = (days, history)
    return history


def _load_finra_daily_ratios(date_value: dt.date) -> dict[str, float]:
    cache_key = date_value.isoformat()
    if cache_key in _finra_daily_cache:
        return _finra_daily_cache[cache_key]

    url = _FINRA_DAILY_URL_TEMPLATE.format(yyyymmdd=date_value.strftime("%Y%m%d"))
    try:
        resp = _http.get(url, timeout=10)
    except requests.RequestException:
        _finra_daily_cache[cache_key] = {}
        return {}

    if resp.status_code != 200:
        _finra_daily_cache[cache_key] = {}
        return {}

    lines = [line.strip() for line in resp.text.splitlines() if line.strip()]
    if not lines:
        _finra_daily_cache[cache_key] = {}
        return {}

    header = [part.strip() for part in lines[0].split("|")]
    try:
        symbol_idx = header.index("Symbol")
        short_idx = header.index("ShortVolume")
        total_idx = header.index("TotalVolume")
    except ValueError:
        _finra_daily_cache[cache_key] = {}
        return {}

    ratios: dict[str, float] = {}
    for line in lines[1:]:
        parts = [part.strip() for part in line.split("|")]
        if len(parts) <= max(symbol_idx, short_idx, total_idx):
            continue
        symbol = parts[symbol_idx].upper()
        try:
            total_volume = float(parts[total_idx])
            short_volume = float(parts[short_idx])
        except ValueError:
            continue
        if total_volume <= 0:
            continue
        ratios[symbol] = short_volume / total_volume

    _finra_daily_cache[cache_key] = ratios
    return ratios


def _fetch_short_ratio_history_from_finra_cdn(ticker: str, event_date: str) -> list[dict]:
    target = ticker.upper()
    event_dt = dt.date.fromisoformat(event_date)
    history: list[dict] = []
    current = event_dt - dt.timedelta(days=1)

    for _ in range(_FINRA_DAILY_MAX_CALENDAR_LOOKBACK_DAYS):
        ratios = _load_finra_daily_ratios(current)
        ratio = ratios.get(target)
        if ratio is not None:
            history.append({"date": current.isoformat(), "short_ratio": ratio})
            if len(history) >= 5:
                break
        current -= dt.timedelta(days=1)

    return history


def fetch_bars(ticker: str, event_date: str) -> list[PriceBar]:
    from datetime import datetime, timedelta
    end_dt = datetime.strptime(event_date, "%Y-%m-%d")
    start_dt = end_dt - timedelta(days=120)
    start_str = start_dt.strftime("%Y-%m-%d")
    # Use full-history cache to avoid per-row HTTP calls
    all_bars_raw = _fetch_full_bars_raw(ticker)
    bars_raw = [b for b in all_bars_raw if start_str <= b.get("date", "") <= event_date]
    bars = []
    for b in bars_raw:
        try:
            bars.append(PriceBar(date=b["date"], open=float(b.get("open", 0)),
                                high=float(b.get("high", 0)), low=float(b.get("low", 0)),
                                close=float(b.get("close", 0)), volume=int(b.get("volume", 0))))
        except (KeyError, ValueError, TypeError):
            continue
    return bars


def _prefetch_ticker_bars(tickers: list, max_workers: int = 16) -> None:
    """Pre-warm full-history bar cache for all unique tickers in parallel."""
    unique = sorted({str(t) for t in tickers if t})
    print(f"  Prefetching price bars for {len(unique)} tickers...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
        for loaded, _ in enumerate(pool.map(_fetch_full_bars_raw, unique), start=1):
            if loaded % 100 == 0 or loaded == len(unique):
                print(f"  bars prefetch {loaded}/{len(unique)}")


async def _prefetch_short_ratio_db_batch(tickers: list[str]) -> None:
    """Bulk-fetch short ratio history for all tickers in a single DB query."""
    dsn = get_settings().postgres_dsn.replace("+asyncpg", "")
    conn = await asyncpg.connect(dsn=dsn)
    try:
        rows = await conn.fetch(
            """
            SELECT ticker_raw,
                   trade_date::text AS date,
                   short_volume::double precision / NULLIF(total_volume::double precision, 0.0) AS short_ratio
            FROM short_sale_daily
            WHERE ticker_raw = ANY($1)
              AND total_volume IS NOT NULL
              AND total_volume > 0
            ORDER BY trade_date DESC
            """,
            tickers,
        )
    finally:
        await conn.close()
    grouped: dict[str, list[dict]] = defaultdict(list)
    for row in rows:
        if row["short_ratio"] is not None:
            grouped[row["ticker_raw"]].append({"date": row["date"], "short_ratio": float(row["short_ratio"])})
    for ticker in tickers:
        _short_ratio_db_cache[ticker] = grouped.get(ticker, [])


def fetch_short_ratio(ticker: str, event_date: str) -> float | None:
    """Fetch average short ratio over 5 days before event."""
    points = _fetch_short_ratio_history_from_db(ticker)
    if not points:
        points = _fetch_short_ratio_history_from_finra_cdn(ticker, event_date)
    if not points:
        points = _fetch_short_ratio_history_from_oracle(
            ticker,
            _required_short_ratio_days(event_date),
        )
    if not points:
        return None
    # Filter to dates before event
    before = []
    for point in points:
        value = point.get("short_ratio")
        if value is None:
            value = point.get("short_percent")
        if point.get("date", "") >= event_date or value is None:
            continue
        try:
            before.append({"date": point["date"], "short_ratio": float(value)})
        except (KeyError, TypeError, ValueError):
            continue
    before.sort(key=lambda p: p["date"], reverse=True)
    recent = before[:5]
    if not recent:
        return None
    return sum(p["short_ratio"] for p in recent) / len(recent)


def fetch_sector_momentum(sector: str, event_date: str) -> float | None:
    """Compute sector ETF 20d return minus SPY 20d return."""
    normalized_sector = normalize_sector_name(sector)
    etf = SECTOR_ETF.get(normalized_sector or "")
    if not etf:
        return None

    etf_bars = _fetch_full_bars_raw(etf)
    etf_dated = {b["date"]: float(b.get("close", 0)) for b in etf_bars if b.get("close")}
    etf_dates = sorted(etf_dated.keys())

    if event_date not in etf_dated or len(etf_dates) < 21:
        return None

    idx = etf_dates.index(event_date)
    if idx < 20:
        return None

    etf_ret = (etf_dated[etf_dates[idx]] - etf_dated[etf_dates[idx - 20]]) / etf_dated[etf_dates[idx - 20]]

    spy_bars = _fetch_full_bars_raw("SPY")
    spy_dated = {b["date"]: float(b.get("close", 0)) for b in spy_bars if b.get("close")}
    spy_dates = sorted(spy_dated.keys())

    if event_date not in spy_dated:
        return None
    spy_idx = spy_dates.index(event_date)
    if spy_idx < 20:
        return None

    spy_ret = (spy_dated[spy_dates[spy_idx]] - spy_dated[spy_dates[spy_idx - 20]]) / spy_dated[spy_dates[spy_idx - 20]]

    return etf_ret - spy_ret


def compute_missing_features(
    ticker: str,
    event_date: str,
    sector: str,
    existing: dict[str, float | None],
) -> dict[str, float | None]:
    result = dict(existing)

    need_price_bars = (
        result.get("pre_event_hurst_60d") is None
        or result.get("pre_event_entropy_60d") is None
    )
    bars = fetch_bars(ticker, event_date) if need_price_bars else []

    if result.get("pre_event_hurst_60d") is None:
        result["pre_event_hurst_60d"] = pre_event_hurst(bars, event_date, 60) if bars else None
    if result.get("pre_event_entropy_60d") is None:
        result["pre_event_entropy_60d"] = pre_event_entropy(bars, event_date, 60) if bars else None

    if result.get("pre_event_short_ratio") is None:
        try:
            result["pre_event_short_ratio"] = fetch_short_ratio(ticker, event_date)
        except Exception:
            result["pre_event_short_ratio"] = None

    if result.get("pre_event_sector_momentum_20d") is None:
        try:
            result["pre_event_sector_momentum_20d"] = fetch_sector_momentum(sector, event_date)
        except Exception:
            result["pre_event_sector_momentum_20d"] = None

    return result


def enrich_split(input_path: Path, output_path: Path):
    table = pq.read_table(input_path)
    n = len(table)
    if n == 0 or "ticker" not in table.column_names or "event_date" not in table.column_names:
        print("  Empty or schema-less split, copying")
        pq.write_table(table, output_path)
        return

    existing_cols = set(table.column_names)
    if all(f in existing_cols for f in FEATURES):
        null_counts = {f: table.column(f).null_count for f in FEATURES}
        if all(v < n * 0.3 for v in null_counts.values()):
            print(f"  Already enriched ({null_counts}), copying")
            pq.write_table(table, output_path)
            return

    tickers = table.column("ticker").to_pylist()
    event_dates = table.column("event_date").to_pylist()
    existing_values = {
        feature: table.column(feature).to_pylist() if feature in existing_cols else [None] * n
        for feature in FEATURES
    }

    # Get sector info if available
    sectors = table.column("asset_type_proxy").to_pylist() if "asset_type_proxy" in table.column_names else ["UNKNOWN"] * n
    # Actually sector might be elsewhere — check for common sector columns
    for col_name in ["sector", "asset_type_proxy"]:
        if col_name in table.column_names:
            sectors = table.column(col_name).to_pylist()
            break

    results = {f: [None] * n for f in FEATURES}
    success = 0

    # Pre-fetch price bars and short ratio history for all unique tickers in batch
    _prefetch_ticker_bars(tickers)
    unique_tickers = sorted({str(t) for t in tickers if t})
    try:
        asyncio.run(_prefetch_short_ratio_db_batch(unique_tickers))
        print(f"  Short ratio DB prefetch done for {len(unique_tickers)} tickers")
    except Exception as e:
        print(f"  Short ratio DB prefetch failed (will fall back per-row): {e}")

    for i in range(n):
        ticker = tickers[i]
        event_date = str(event_dates[i])
        sector = str(sectors[i]) if sectors[i] else "UNKNOWN"

        if i % 100 == 0:
            print(f"  {i}/{n} ({success} ok)...")

        existing = {feature: existing_values[feature][i] for feature in FEATURES}
        if all(value is not None for value in existing.values()):
            for feature in FEATURES:
                results[feature][i] = existing[feature]
            success += 1
            continue

        try:
            feats = compute_missing_features(ticker, event_date, sector, existing)
            for f in FEATURES:
                results[f][i] = feats.get(f)
            if feats.get("pre_event_hurst_60d") is not None:
                success += 1
        except Exception as e:
            if i < 5:
                print(f"  Error {ticker} {event_date}: {e}")

    print(f"  Done: {success}/{n}")

    for f in FEATURES:
        arr = pa.array(results[f], type=pa.float64())
        if f in existing_cols:
            idx = table.column_names.index(f)
            table = table.set_column(idx, f, arr)
        else:
            table = table.append_column(f, arr)

    pq.write_table(table, output_path)
    print(f"  Written to {output_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    input_dir = Path(args.input)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    manifest_src = input_dir / "manifest.json"
    if manifest_src.exists():
        manifest = json.loads(manifest_src.read_text())
        source_snapshot_id = manifest.get("snapshot_id")
        manifest["snapshot_id"] = output_dir.name
        manifest["output_dir"] = str(output_dir.resolve())
        manifest["created_at_utc"] = dt.datetime.now(dt.UTC).isoformat()
        manifest["enrichment_source_snapshot_id"] = source_snapshot_id
        manifest["export_enrichments"] = sorted(set((manifest.get("export_enrichments") or []) + FEATURES))
        (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))

    for split in ["train", "valid", "test"]:
        input_path = input_dir / f"{split}.parquet"
        if not input_path.exists():
            continue
        output_path = output_dir / f"{split}.parquet"
        print(f"Enriching {split}...")
        enrich_split(input_path, output_path)

    print("All done!")


if __name__ == "__main__":
    main()