"""Compare high-down 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
from typing import Iterable


ANCHORS = {
    "x150_y61": {"x": (148, 152), "y": (57, 65), "dy": 8},
    "x158_y77": {"x": (156, 160), "y": (73, 81), "dy": 8},
    "x166_y93": {"x": (164, 168), "y": (89, 97), "dy": 8},
    "x174_y109": {"x": (172, 176), "y": (105, 113), "dy": 8},
    "x182_y125": {"x": (180, 184), "y": (121, 129), "dy": 8},
    "x190_y141": {"x": (188, 192), "y": (137, 145), "dy": 8},
    "x198_y157": {"x": (196, 200), "y": (153, 161), "dy": 8},
}


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


def matches_anchor(row: dict, spec: dict) -> bool:
    ball_x = row.get("ball_x")
    ball_y = row.get("ball_y")
    return (
        row.get("dx") is not None
        and row.get("dx") > 0
        and row.get("dy") == spec["dy"]
        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]
    )


def compact_row(row: dict) -> dict:
    target_y = row.get("target_y")
    paddle_y = row.get("self_paddle_y")
    paddle_error = None
    if target_y is not None and paddle_y is not None:
        paddle_error = round(float(paddle_y) - float(target_y), 2)
    return {
        "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"),
        "dx": row.get("dx"),
        "dy": row.get("dy"),
        "target_y": row.get("target_y"),
        "self_paddle_y": row.get("self_paddle_y"),
        "paddle_error": paddle_error,
        "action_name": row.get("action_name"),
        "mode": row.get("mode"),
    }


def summarize_rows(rows: list[dict]) -> dict:
    if not rows:
        return {"count": 0}

    def bucket(row: dict) -> tuple:
        return (
            row.get("ball_x"),
            row.get("ball_y"),
            row.get("target_y"),
            row.get("self_paddle_y"),
            row.get("action_name"),
            row.get("mode"),
            f"{row.get('self_score')}-{row.get('opponent_score')}",
        )

    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(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],
            "target_y": top_bucket[2],
            "self_paddle_y": top_bucket[3],
            "action_name": top_bucket[4],
            "mode": top_bucket[5],
            "score": top_bucket[6],
        },
        "mean_paddle_error": round(mean(errors), 3) if errors else None,
        "action_counts": dict(Counter(row.get("action_name") for row in rows)),
        "score_counts": dict(Counter(f"{row.get('self_score')}-{row.get('opponent_score')}" for row in rows)),
    }


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

    for record in read_records(path):
        window = record.get("window") or []
        outcome_key = record.get("pattern") if record.get("outcome") == "loss" else "win"
        for name, spec in ANCHORS.items():
            for row in window:
                if not matches_anchor(row, spec):
                    continue
                anchor_rows[name][outcome_key].append(row)
                final_patterns[name][outcome_key] += 1
                examples[name].setdefault(
                    outcome_key,
                    {
                        "seed": record.get("seed"),
                        "point_index": record.get("point_index"),
                        "record_outcome": record.get("outcome"),
                        "pattern": record.get("pattern"),
                        "anchor": compact_row(row),
                        "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"),
                        },
                    },
                )
                break

    anchors = {}
    for name in ANCHORS:
        outcomes = {
            outcome: summarize_rows(rows)
            for outcome, rows in sorted(anchor_rows[name].items())
        }
        anchors[name] = {
            "outcomes": outcomes,
            "outcome_counts": dict(final_patterns[name]),
            "examples": examples.get(name, {}),
        }

    return {"source": str(path), "anchors": 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()
