"""Summarize seed-level deltas and intervention contexts for counterfactual probes."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path
from statistics import mean


def read_rows(path: Path) -> list[dict]:
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            rows.append(json.loads(line))
    return rows


def first_intervention_bucket(row: dict) -> str:
    interventions = row.get("interventions") or []
    if not interventions:
        return "none"
    first = interventions[0]
    return (
        f"score={first.get('self_score')}-{first.get('opponent_score')};"
        f"x={first.get('ball_x')};y={first.get('ball_y')};"
        f"pad={first.get('self_paddle_y')};"
        f"target={round(float(first.get('target_y')), 2)};"
        f"action={first.get('baseline_action')}->{first.get('override_action')};"
        f"step={first.get('step')}"
    )


def summarize(path: Path, baseline_variant: str) -> dict:
    rows = read_rows(path)
    baseline = {row["seed"]: row for row in rows if row["variant"] == baseline_variant}
    variants = sorted({row["variant"] for row in rows if row["variant"] != baseline_variant})

    result = {"source": str(path), "baseline_variant": baseline_variant, "variants": {}}
    for variant in variants:
        variant_rows = [row for row in rows if row["variant"] == variant]
        deltas = []
        improved = []
        worsened = []
        unchanged = []
        first_buckets = Counter()
        first_buckets_by_delta = {"improved": Counter(), "worsened": Counter(), "unchanged": Counter()}

        for row in variant_rows:
            base = baseline[row["seed"]]
            delta = float(row["score"]) - float(base["score"])
            deltas.append(delta)
            seed_summary = {
                "seed": row["seed"],
                "baseline": base["score"],
                "score": row["score"],
                "delta": delta,
                "intervention_count": row.get("intervention_count", 0),
                "first_intervention": (row.get("interventions") or [None])[0],
            }
            bucket = first_intervention_bucket(row)
            first_buckets[bucket] += 1
            if delta > 0:
                improved.append(seed_summary)
                first_buckets_by_delta["improved"][bucket] += 1
            elif delta < 0:
                worsened.append(seed_summary)
                first_buckets_by_delta["worsened"][bucket] += 1
            else:
                unchanged.append(seed_summary)
                first_buckets_by_delta["unchanged"][bucket] += 1

        result["variants"][variant] = {
            "mean_delta": round(mean(deltas), 4) if deltas else None,
            "min_delta": min(deltas) if deltas else None,
            "max_delta": max(deltas) if deltas else None,
            "improved_count": len(improved),
            "worsened_count": len(worsened),
            "unchanged_count": len(unchanged),
            "top_first_intervention_buckets": dict(first_buckets.most_common(12)),
            "top_first_intervention_buckets_by_delta": {
                key: dict(counter.most_common(8))
                for key, counter in first_buckets_by_delta.items()
            },
            "largest_improvements": sorted(improved, key=lambda item: item["delta"], reverse=True)[:10],
            "largest_regressions": sorted(worsened, key=lambda item: item["delta"])[:20],
        }
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("jsonl", type=Path)
    parser.add_argument("--baseline-variant", default="baseline")
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    result = summarize(args.jsonl, args.baseline_variant)
    text = json.dumps(result, indent=2, sort_keys=True)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()
