"""Compare compact pre-anchor histories for high-down trajectories."""

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


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 is_x150_high_down(row: 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") == 8
        and ball_x is not None
        and 148 <= ball_x <= 152
        and ball_y is not None
        and 57 <= ball_y <= 65
    )


def outcome_key(record: dict) -> str | None:
    if record.get("outcome") == "win":
        return "win"
    if record.get("pattern") == "loss/high_down/paddle_above_target":
        return "loss_high_down"
    return None


def paddle_error(row: dict) -> float | None:
    target_y = row.get("target_y")
    paddle_y = row.get("self_paddle_y")
    if target_y is None or paddle_y is None:
        return None
    return float(paddle_y) - float(target_y)


def error_bucket(value: float | None) -> str:
    if value is None:
        return "na"
    rounded = int(round(value / 5.0) * 5)
    return str(rounded)


def compact_sample(row: dict) -> dict:
    error = paddle_error(row)
    return {
        "step": row.get("step"),
        "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": round(error, 2) if error is not None else None,
        "action_name": row.get("action_name"),
        "mode": row.get("mode"),
        "score": f"{row.get('self_score')}-{row.get('opponent_score')}",
    }


def history_signature(rows: list[dict]) -> str:
    parts = []
    for row in rows:
        action = (row.get("action_name") or "?")[:1]
        mode = (row.get("mode") or "?")[:1]
        parts.append(f"{int(row.get('ball_x') or -1)}:{action}{mode}:{error_bucket(paddle_error(row))}")
    return "|".join(parts)


def summarize_offset(rows: list[dict]) -> dict:
    if not rows:
        return {"count": 0}
    errors = [paddle_error(row) for row in rows]
    errors = [value for value in errors if value is not None]
    return {
        "count": len(rows),
        "action_counts": dict(Counter(row.get("action_name") for row in rows)),
        "mode_counts": dict(Counter(row.get("mode") for row in rows)),
        "target_counts": dict(Counter(row.get("target_y") for row in rows)),
        "paddle_counts": dict(Counter(row.get("self_paddle_y") for row in rows)),
        "mean_paddle_error": round(mean(errors), 3) if errors else None,
    }


def compare(path: Path, history: int) -> dict:
    by_outcome: dict[str, list[dict]] = defaultdict(list)
    offset_rows: dict[str, dict[int, list[dict]]] = defaultdict(lambda: defaultdict(list))
    signatures: dict[str, Counter] = defaultdict(Counter)
    examples: dict[str, dict] = {}

    for record in read_records(path):
        key = outcome_key(record)
        if key is None:
            continue
        window = record.get("window") or []
        anchor_index = next((i for i, row in enumerate(window) if is_x150_high_down(row)), None)
        if anchor_index is None:
            continue

        start = max(0, anchor_index - history)
        rows = window[start : anchor_index + 1]
        by_outcome[key].append(
            {
                "seed": record.get("seed"),
                "point_index": record.get("point_index"),
                "record_step": record.get("step"),
                "pattern": record.get("pattern"),
                "anchor": compact_sample(window[anchor_index]),
                "history": [compact_sample(row) for row in rows],
                "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"),
                },
            }
        )
        signatures[key][history_signature(rows)] += 1
        examples.setdefault(key, by_outcome[key][-1])

        for rel in range(-history, 1):
            idx = anchor_index + rel
            if 0 <= idx < len(window):
                offset_rows[key][rel].append(window[idx])

    outcomes = {}
    for key, records in sorted(by_outcome.items()):
        outcomes[key] = {
            "count": len(records),
            "top_history_signatures": [
                {"count": count, "signature": signature}
                for signature, count in signatures[key].most_common(10)
            ],
            "offsets": {
                str(rel): summarize_offset(rows)
                for rel, rows in sorted(offset_rows[key].items())
            },
            "example": examples.get(key),
        }

    return {
        "source": str(path),
        "anchor": "x150_high_down",
        "history": history,
        "outcomes": outcomes,
    }


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

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