fithia2/scripts/enrich_technical_features.py

"""Enrich snapshot with technical indicator features using PyArrow (no pandas round-trip).

Computes: pre_event_volatility_20d, pre_event_rsi_14, pre_event_bb_position, pre_event_obv_slope_20d
from Oracle price bars for each event in the snapshot.

Usage:
    uv run python3 scripts/enrich_technical_features.py \
        --input data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_mom2 \
        --output data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_tech
"""
from __future__ import annotations

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

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

# Add project root to path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from libs.features.market_features import (
    pre_event_bb_position,
    pre_event_obv_slope,
    pre_event_rsi,
    pre_event_volatility,
)
from libs.oracle_client.models import PriceBar

ORACLE_URL = "http://localhost:18001"
FEATURES = ["pre_event_volatility_20d", "pre_event_rsi_14", "pre_event_bb_position", "pre_event_obv_slope_20d"]


def fetch_bars(ticker: str, event_date: str) -> list[PriceBar]:
    """Fetch 60 days of bars ending on event_date from Oracle."""
    from datetime import datetime, timedelta
    end_dt = datetime.strptime(event_date, "%Y-%m-%d")
    start_dt = end_dt - timedelta(days=90)  # 90 calendar days ≈ 60+ trading days
    url = f"{ORACLE_URL}/api/v1/price/data/{ticker}"
    resp = requests.get(url, params={
        "start_date": start_dt.strftime("%Y-%m-%d"),
        "end_date": event_date,
    }, timeout=30)
    if resp.status_code != 200:
        return []
    data = resp.json()
    bars_raw = data.get("bars") or data.get("data") or []
    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 compute_features_for_event(ticker: str, event_date: str) -> dict[str, float | None]:
    bars = fetch_bars(ticker, event_date)
    if not bars:
        return {f: None for f in FEATURES}
    return {
        "pre_event_volatility_20d": pre_event_volatility(bars, event_date, 20),
        "pre_event_rsi_14": pre_event_rsi(bars, event_date, 14),
        "pre_event_bb_position": pre_event_bb_position(bars, event_date, 20),
        "pre_event_obv_slope_20d": pre_event_obv_slope(bars, event_date, 20),
    }


def enrich_split(input_path: Path, output_path: Path):
    table = pq.read_table(input_path)
    n = len(table)

    # Check if already enriched
    existing_cols = set(table.column_names)
    if all(f in existing_cols for f in FEATURES):
        # Check null counts
        null_counts = {f: table.column(f).null_count for f in FEATURES}
        if all(v < n * 0.1 for v in null_counts.values()):
            print(f"  Already enriched ({null_counts}), copying as-is")
            pq.write_table(table, output_path)
            return

    tickers = table.column("ticker").to_pylist()
    event_dates = table.column("event_date").to_pylist()

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

    for i in range(n):
        ticker = tickers[i]
        event_date = str(event_dates[i])
        if i % 100 == 0:
            print(f"  Processing {i}/{n} ({success} ok, {errors} err)...")

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

        # No sleep needed — sequential Oracle calls are fine

    print(f"  Done: {success} success, {errors} errors out of {n}")

    # Append columns using PyArrow (no pandas)
    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, help="Input snapshot dir")
    parser.add_argument("--output", required=True, help="Output snapshot dir")
    args = parser.parse_args()

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

    # Copy manifest
    manifest_src = input_dir / "manifest.json"
    if manifest_src.exists():
        manifest = json.loads(manifest_src.read_text())
        manifest["snapshot_id"] = output_dir.name
        (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():
            print(f"Skipping {split} (not found)")
            continue
        output_path = output_dir / f"{split}.parquet"
        print(f"Enriching {split}...")
        enrich_split(input_path, output_path)

    print("Done!")


if __name__ == "__main__":
    main()