"""Enrich snapshot with prior_event_fwd5d (cross-event momentum feature).

Computes from existing snapshot data — no Oracle API needed.
PIT-safe: only uses realized forward returns from prior events.

Usage:
    uv run python3 scripts/enrich_prior_drift.py \
        --input data/parquet/midwide-liquid-long-v1_tier3_macro \
        --output data/parquet/midwide-liquid-long-v1_full
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import sys
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

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


def trading_days_between(start: dt.date, end: dt.date) -> list[dt.date]:
    """Count approximate trading days (Mon-Fri, no holiday exclusion)."""
    days = []
    cur = start
    while cur <= end:
        if cur.weekday() < 5:
            days.append(cur)
        cur += dt.timedelta(days=1)
    return days


def compute_prior_drift(rows: list[dict]) -> list[float | None]:
    """Compute prior_event_fwd5d for all rows. PIT-safe."""
    sorted_rows = sorted(
        rows,
        key=lambda r: (
            str(r.get("ticker", "")),
            str(r.get("event_date", "")),
            str(r.get("entry_date", "")),
        ),
    )

    realized_on_cache: dict[dt.date, dt.date | None] = {}

    def _fwd5_realized_on(entry_date_str: str | None) -> dt.date | None:
        if not entry_date_str:
            return None
        try:
            entry_date = dt.date.fromisoformat(str(entry_date_str))
        except (ValueError, TypeError):
            return None
        cached = realized_on_cache.get(entry_date)
        if cached is not None or entry_date in realized_on_cache:
            return cached
        tdays = trading_days_between(entry_date, entry_date + dt.timedelta(days=14))
        realized_on = tdays[5] if len(tdays) > 5 else None
        realized_on_cache[entry_date] = realized_on
        return realized_on

    prev_by_ticker: dict[str, tuple[float | None, dt.date | None]] = {}
    result_map: dict[int, float | None] = {}

    for row in sorted_rows:
        orig_idx = row["_orig_idx"]
        ticker = str(row.get("ticker", ""))
        try:
            current_event_date = dt.date.fromisoformat(str(row.get("event_date", "")))
        except (ValueError, TypeError):
            result_map[orig_idx] = None
            continue

        prior_value: float | None = None
        prior = prev_by_ticker.get(ticker)
        if prior is not None:
            candidate_value, realized_on = prior
            if realized_on is not None and current_event_date > realized_on:
                prior_value = candidate_value

        result_map[orig_idx] = prior_value

        fwd5 = row.get("fwd_return_5d")
        entry_date_str = row.get("entry_date")
        prev_by_ticker[ticker] = (
            float(fwd5) if fwd5 is not None else None,
            _fwd5_realized_on(str(entry_date_str) if entry_date_str else None),
        )

    return [result_map[i] for i in range(len(rows))]


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

    if "prior_event_fwd5d" in table.column_names:
        null_count = table.column("prior_event_fwd5d").null_count
        if null_count < n * 0.3:
            print(f"  Already has prior_event_fwd5d ({null_count} nulls), copying as-is")
            pq.write_table(table, output_path)
            return

    cols = table.column_names
    rows = []
    tickers = table.column("ticker").to_pylist() if "ticker" in cols else [None] * n
    event_dates = table.column("event_date").to_pylist() if "event_date" in cols else [None] * n
    entry_dates = table.column("entry_date").to_pylist() if "entry_date" in cols else [None] * n
    fwd5s = table.column("fwd_return_5d").to_pylist() if "fwd_return_5d" in cols else [None] * n

    for i in range(n):
        rows.append({
            "_orig_idx": i,
            "ticker": tickers[i],
            "event_date": str(event_dates[i]) if event_dates[i] else None,
            "entry_date": str(entry_dates[i]) if entry_dates[i] else None,
            "fwd_return_5d": fwd5s[i],
        })

    values = compute_prior_drift(rows)
    enriched = sum(1 for v in values if v is not None)
    print(f"  prior_event_fwd5d: {enriched}/{n} non-null")

    arr = pa.array(values, type=pa.float64())
    if "prior_event_fwd5d" in table.column_names:
        idx = table.column_names.index("prior_event_fwd5d")
        table = table.set_column(idx, "prior_event_fwd5d", arr)
    else:
        table = table.append_column("prior_event_fwd5d", 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():
        import json as _json
        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"\nProcessing {split}...")
        enrich_split(input_path, output_path)

    print("\nDone.")


if __name__ == "__main__":
    main()