fithia2/scripts/enrich_earnings_history_features.py

"""Enrich snapshot with historical earnings surprise features.

Adds PIT-safe trailing quarterly surprise features derived from Oracle's
earnings surprise endpoint. The current event's earnings surprise is already
materialized separately as `earnings_surprise_pct`; this script adds the
prior 12 quarter surprise vector plus compact summary statistics.

Usage:
    uv run python3 scripts/enrich_earnings_history_features.py \
        --input data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical \
        --output data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical_eh
"""
from __future__ import annotations

import argparse
import concurrent.futures
import datetime as dt
import json
import sys
from pathlib import Path
from typing import Any

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

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

ORACLE_URL = "http://localhost:18001"
SURPRISE_VECTOR_COLS = [f"sue_lag_{i}_pct" for i in range(1, 13)]
SUMMARY_COLS = [
    "sue_hist_mean_4q",
    "sue_hist_mean_8q",
    "sue_hist_mean_12q",
    "sue_hist_pos_rate_4q",
    "sue_hist_pos_rate_12q",
    "sue_hist_latest_pct",
    "sue_hist_streak_pos",
]
FEATURES = SURPRISE_VECTOR_COLS + SUMMARY_COLS


def _parse_date(value: Any) -> dt.date | None:
    if value is None:
        return None
    text = str(value)
    if not text:
        return None
    try:
        return dt.date.fromisoformat(text[:10])
    except ValueError:
        return None


def _fetch_surprise_quarters(ticker: str) -> list[dict[str, Any]]:
    resp = requests.get(f"{ORACLE_URL}/api/v1/earnings/surprise/{ticker}", timeout=30)
    if resp.status_code != 200:
        return []
    data = resp.json()
    quarters = data.get("quarters") or []
    return [q for q in quarters if q.get("reported_date")]


def _match_current_quarter(quarters: list[dict[str, Any]], event_date: dt.date) -> int | None:
    best_idx: int | None = None
    best_delta = 999999
    for idx, quarter in enumerate(quarters):
        reported_date = _parse_date(quarter.get("reported_date"))
        if reported_date is None:
            continue
        delta = abs((reported_date - event_date).days)
        if delta <= 5 and delta < best_delta:
            best_delta = delta
            best_idx = idx
    return best_idx


def _mean(values: list[float]) -> float | None:
    if not values:
        return None
    return sum(values) / len(values)


def _positive_rate(values: list[float]) -> float | None:
    if not values:
        return None
    return sum(1 for value in values if value > 0.0) / len(values)


def _positive_streak(values: list[float]) -> float | None:
    if not values:
        return None
    streak = 0
    for value in values:
        if value > 0.0:
            streak += 1
        else:
            break
    return float(streak)


def compute_history_features(
    ticker: str,
    event_date: str | dt.date | None,
    quarters: list[dict[str, Any]] | None = None,
) -> dict[str, float | None]:
    event_dt = _parse_date(event_date)
    result = {feature: None for feature in FEATURES}
    if event_dt is None:
        return result

    if quarters is None:
        quarters = _fetch_surprise_quarters(ticker)
    if not quarters:
        return result

    matched_idx = _match_current_quarter(quarters, event_dt)
    if matched_idx is None:
        return result

    prior_quarters = quarters[matched_idx + 1 : matched_idx + 13]
    prior_values: list[float] = []
    for idx, quarter in enumerate(prior_quarters, start=1):
        surprise_pct = quarter.get("surprise_percentage")
        if surprise_pct is None:
            prior_values.append(float("nan"))
            continue
        value = float(surprise_pct)
        prior_values.append(value)
        result[f"sue_lag_{idx}_pct"] = value

    clean_values = [value for value in prior_values if value == value]
    clean_4 = clean_values[:4]
    clean_8 = clean_values[:8]
    clean_12 = clean_values[:12]

    result["sue_hist_mean_4q"] = _mean(clean_4)
    result["sue_hist_mean_8q"] = _mean(clean_8)
    result["sue_hist_mean_12q"] = _mean(clean_12)
    result["sue_hist_pos_rate_4q"] = _positive_rate(clean_4)
    result["sue_hist_pos_rate_12q"] = _positive_rate(clean_12)
    result["sue_hist_latest_pct"] = clean_values[0] if clean_values else None
    result["sue_hist_streak_pos"] = _positive_streak(clean_values)
    return result


def enrich_split(input_path: Path, output_path: Path) -> None:
    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(feature in existing_cols for feature in FEATURES):
        null_counts = {feature: table.column(feature).null_count for feature in FEATURES}
        if all(count < n * 0.5 for count 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()
    event_types = table.column("event_type").to_pylist() if "event_type" in table.column_names else [None] * n

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

    earnings_tickers = sorted(
        {
            str(tickers[idx])
            for idx in range(n)
            if event_types[idx] == "earnings_release" and tickers[idx]
        }
    )
    surprise_cache: dict[str, list[dict[str, Any]]] = {}
    if earnings_tickers:
        print(f"  Prefetching surprise history for {len(earnings_tickers)} tickers...")

        def _load_one(ticker: str) -> tuple[str, list[dict[str, Any]]]:
            try:
                return ticker, _fetch_surprise_quarters(ticker)
            except Exception:
                return ticker, []

        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
            for loaded, (ticker, quarters) in enumerate(pool.map(_load_one, earnings_tickers), start=1):
                surprise_cache[ticker] = quarters
                if loaded % 100 == 0 or loaded == len(earnings_tickers):
                    print(f"  surprise prefetch {loaded}/{len(earnings_tickers)}")

    for idx in range(n):
        ticker = str(tickers[idx]) if tickers[idx] is not None else ""
        event_date = event_dates[idx]
        event_type = str(event_types[idx] or "")

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

        if event_type != "earnings_release" or not ticker:
            continue

        quarters = surprise_cache.get(ticker, [])

        try:
            features = compute_history_features(ticker, event_date, quarters)
            if any(features[feature] is not None for feature in FEATURES):
                success += 1
            for feature in FEATURES:
                results[feature][idx] = features[feature]
        except Exception as exc:
            if idx < 5:
                print(f"  Error {ticker} {event_date}: {exc}")

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

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

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


def main() -> None:
    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()