"""Compare y86 trajectory anchors across point outcomes."""

from __future__ import annotations

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


ANCHORS = {
    "x126_y81_up": {"x": (124, 128), "y": (79, 83), "dy": (-4, -4)},
    "x142_y65_up": {"x": (140, 144), "y": (63, 67), "dy": (-4, -4)},
    "x150_y57_up": {"x": (148, 152), "y": (55, 59), "dy": (-4, -4)},
    "x158_y49_up": {"x": (156, 160), "y": (47, 51), "dy": (-4, -4)},
    "x166_y50_bounce": {"x": (164, 168), "y": (48, 52), "dy": (3, 5)},
    "x174_y58_down": {"x": (172, 176), "y": (56, 60), "dy": (4, 4)},
    "x182_y66_down": {"x": (180, 184), "y": (64, 68), "dy": (4, 4)},
    "x190_y74_down": {"x": (188, 192), "y": (72, 76), "dy": (4, 4)},
    "x194_y78_down": {"x": (192, 196), "y": (76, 80), "dy": (4, 4)},
    "x198_y82_down": {"x": (196, 200), "y": (80, 84), "dy": (4, 4)},
    "x202_y86_down": {"x": (200, 204), "y": (84, 88), "dy": (4, 4)},
}


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


def outcome_key(record: dict) -> str:
    if record.get("outcome") == "win":
        return "win"
    return str(record.get("pattern"))


def matches(row: dict, spec: dict) -> bool:
    ball_x = row.get("ball_x")
    ball_y = row.get("ball_y")
    dy = row.get("dy")
    return (
        row.get("dx") is not None
        and row.get("dx") > 0
        and ball_x is not None
        and spec["x"][0] <= ball_x <= spec["x"][1]
        and ball_y is not None
        and spec["y"][0] <= ball_y <= spec["y"][1]
        and dy is not None
        and spec["dy"][0] <= dy <= spec["dy"][1]
    )


def row_bucket(row: dict) -> tuple:
    return (
        row.get("ball_x"),
        row.get("ball_y"),
        row.get("dy"),
        row.get("target_y"),
        row.get("self_paddle_y"),
        row.get("action_name"),
        row.get("mode"),
        row.get("self_score"),
        row.get("opponent_score"),
    )


def summarize_rows(rows: list[dict]) -> dict:
    if not rows:
        return {"count": 0}
    errors = [
        float(row["self_paddle_y"]) - float(row["target_y"])
        for row in rows
        if row.get("self_paddle_y") is not None and row.get("target_y") is not None
    ]
    top_bucket, top_count = Counter(row_bucket(row) for row in rows).most_common(1)[0]
    return {
        "count": len(rows),
        "top_count": top_count,
        "top_fraction": round(top_count / len(rows), 3),
        "top": {
            "ball_x": top_bucket[0],
            "ball_y": top_bucket[1],
            "dy": top_bucket[2],
            "target_y": top_bucket[3],
            "self_paddle_y": top_bucket[4],
            "action_name": top_bucket[5],
            "mode": top_bucket[6],
            "score": f"{top_bucket[7]}-{top_bucket[8]}",
        },
        "mean_paddle_error": round(mean(errors), 3) if errors else None,
        "action_counts": dict(Counter(row.get("action_name") for row in rows)),
        "mode_counts": dict(Counter(row.get("mode") for row in rows)),
        "score_counts": dict(Counter(f"{row.get('self_score')}-{row.get('opponent_score')}" for row in rows)),
    }


def compact_example(record: dict, row: dict) -> dict:
    return {
        "seed": record.get("seed"),
        "point_index": record.get("point_index"),
        "outcome": record.get("outcome"),
        "pattern": record.get("pattern"),
        "anchor": {
            "step": row.get("step"),
            "score": f"{row.get('self_score')}-{row.get('opponent_score')}",
            "ball_x": row.get("ball_x"),
            "ball_y": row.get("ball_y"),
            "dy": row.get("dy"),
            "target_y": row.get("target_y"),
            "self_paddle_y": row.get("self_paddle_y"),
            "action_name": row.get("action_name"),
            "mode": row.get("mode"),
        },
        "final": {
            "ball_x": record.get("final_ball_x"),
            "ball_y": record.get("final_ball_y"),
            "dy": record.get("final_dy"),
            "target_y": record.get("final_target_y"),
            "self_paddle_y": record.get("final_self_paddle_y"),
            "paddle_error": record.get("final_paddle_error"),
        },
    }


def compare(path: Path) -> dict:
    records = read_records(path)
    anchor_rows: dict[str, dict[str, list[dict]]] = defaultdict(lambda: defaultdict(list))
    examples: dict[str, dict[str, dict]] = defaultdict(dict)

    for record in records:
        key = outcome_key(record)
        window = record.get("window") or []
        for name, spec in ANCHORS.items():
            for row in window:
                if not matches(row, spec):
                    continue
                anchor_rows[name][key].append(row)
                examples[name].setdefault(key, compact_example(record, row))
                break

    return {
        "source": str(path),
        "anchors": {
            name: {
                "outcomes": {
                    key: summarize_rows(rows)
                    for key, rows in sorted(anchor_rows[name].items())
                },
                "examples": examples.get(name, {}),
            }
            for name in ANCHORS
        },
    }


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

    result = compare(args.jsonl)
    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()
