"""Event Type Decomposition: Per-event-type performance analysis.

Which event types contribute alpha?

Groups snapshot data by event_type and computes forward return statistics,
win rates, and average composite scores per type.

Usage:
  python -m dev.analysis.event_type_decomposition \
    --snapshot-dir data/datasets/snapshots/snapshot_2026_03_20 \
    --split train [--output-dir ./data/analysis]
"""
from __future__ import annotations

import argparse
import csv
import statistics
from pathlib import Path
from typing import Any

import pyarrow.parquet as pq

from libs.backtest.scoring import compute_entry_score


def _load_rows(snapshot_dir: Path, split: str) -> list[dict[str, Any]]:
    parquet_path = snapshot_dir / f"{split}.parquet"
    table = pq.read_table(parquet_path)
    return table.to_pylist()


FWD_HORIZONS = ["fwd_return_1d", "fwd_return_3d", "fwd_return_5d", "fwd_return_10d"]


def analyze(rows: list[dict[str, Any]], output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    # Compute scores
    for row in rows:
        row["_composite_score"] = compute_entry_score(row)

    # Group by event_type
    by_type: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        et = row.get("event_type", "unknown")
        by_type.setdefault(et, []).append(row)

    csv_rows: list[dict[str, Any]] = []
    for event_type in sorted(by_type):
        items = by_type[event_type]
        n = len(items)
        row_out: dict[str, Any] = {"event_type": event_type, "count": n}

        # Average composite score
        scores = [r["_composite_score"] for r in items]
        row_out["avg_score"] = statistics.mean(scores) if scores else None

        for horizon in FWD_HORIZONS:
            vals = [float(r[horizon]) for r in items if r.get(horizon) is not None]
            if vals:
                row_out[f"{horizon}_mean"] = statistics.mean(vals)
                row_out[f"{horizon}_win_rate"] = sum(1 for v in vals if v > 0) / len(vals)
                row_out[f"{horizon}_std"] = statistics.stdev(vals) if len(vals) > 1 else 0.0
            else:
                row_out[f"{horizon}_mean"] = None
                row_out[f"{horizon}_win_rate"] = None
                row_out[f"{horizon}_std"] = None

        csv_rows.append(row_out)

    # Write CSV
    csv_path = output_dir / "event_type_decomposition.csv"
    if csv_rows:
        fieldnames = list(csv_rows[0].keys())
        with open(csv_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(csv_rows)

    # Console output
    print(f"\n{'='*90}")
    print(f"Event Type Decomposition — {len(rows)} events, {len(by_type)} types")
    print(f"{'='*90}")
    print(f"\n{'Type':<22} {'N':>5} {'AvgScore':>8}", end="")
    for h in FWD_HORIZONS:
        label = h.replace("fwd_return_", "")
        print(f"  {label+'_mean':>9} {label+'_wr':>7}", end="")
    print()
    print("-" * 90)

    for row_out in csv_rows:
        avg_s = row_out["avg_score"]
        print(f"{row_out['event_type']:<22} {row_out['count']:>5} {avg_s:.3f}" if avg_s else
              f"{row_out['event_type']:<22} {row_out['count']:>5}     N/A", end="")
        for h in FWD_HORIZONS:
            mean_val = row_out[f"{h}_mean"]
            wr_val = row_out[f"{h}_win_rate"]
            mean_str = f"{mean_val:+.4f}" if mean_val is not None else "     N/A"
            wr_str = f"{wr_val:.1%}" if wr_val is not None else "  N/A"
            print(f"  {mean_str:>9} {wr_str:>7}", end="")
        print()

    print(f"\nCSV written to: {csv_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Event Type Decomposition Analysis")
    parser.add_argument("--snapshot-dir", required=True, help="Path to snapshot directory")
    parser.add_argument("--split", default="train", help="Split name (train/valid/test)")
    parser.add_argument("--output-dir", default="./data/analysis", help="Output directory")
    args = parser.parse_args()

    rows = _load_rows(Path(args.snapshot_dir), args.split)
    if not rows:
        print("No data found.")
        return

    analyze(rows, Path(args.output_dir))


if __name__ == "__main__":
    main()
