"""Compare setup-route windows for selected point-trace route keys."""

from __future__ import annotations

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


Trace = dict[str, Any]


def read_jsonl(path: Path) -> list[Trace]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")


def bucket(value: int | float, width: int) -> str:
    low = int(value) // width * width
    return f"{low}..{low + width - 1}"


def route_key(pre: Trace, entry: Trace, right: Trace) -> str:
    return (
        f"pre_y={bucket(pre['ball_y'], 20)};"
        f"entry_y={bucket(entry['ball_y'], 20)};"
        f"right_y={bucket(right['ball_y'], 20)};"
        f"entry_pad={bucket(entry['self_paddle_y'], 20)};"
        f"right_pad={bucket(right['self_paddle_y'], 20)}"
    )


def first_right_approach_index(window: list[Trace]) -> int | None:
    for index, row in enumerate(window):
        if (
            row.get("dx") is not None
            and row["dx"] > 0
            and row["ball_x"] >= 120
            and row["ball_y"] != 0
            and row["ball_x"] not in (0, 205)
            and row.get("mode") != "inactive"
        ):
            return index
    return None


def route_for_record(record: Trace) -> str | None:
    window = record.get("window") or []
    first_right = first_right_approach_index(window)
    if first_right is None:
        return None
    pre_index = max(0, first_right - 1)
    right_index = min(len(window) - 1, first_right + 3)
    return route_key(window[pre_index], window[first_right], window[right_index])


def compact_row(row: Trace | None) -> dict[str, Any] | None:
    if row is None:
        return None
    return {
        "ball_x": row.get("ball_x"),
        "ball_y": row.get("ball_y"),
        "dx": row.get("dx"),
        "dy": row.get("dy"),
        "target_y": round(float(row["target_y"]), 2),
        "paddle_y": row.get("self_paddle_y"),
        "paddle_error": round(float(row["self_paddle_y"]) - float(row["target_y"]), 2),
        "action": row.get("action_name"),
        "mode": row.get("mode"),
    }


def row_bucket(row: Trace | None) -> str:
    if row is None:
        return "missing"
    return (
        f"bx={row.get('ball_x')};by={row.get('ball_y')};"
        f"dx={row.get('dx')};dy={row.get('dy')};"
        f"target={round(float(row['target_y']), 2)};"
        f"pad={row.get('self_paddle_y')};"
        f"err={round(float(row['self_paddle_y']) - float(row['target_y']), 2)};"
        f"action={row.get('action_name')};mode={row.get('mode')}"
    )


def numeric_summary(values: list[float]) -> dict[str, Any]:
    if not values:
        return {"count": 0, "mean": None, "min": None, "max": None}
    return {"count": len(values), "mean": mean(values), "min": min(values), "max": max(values)}


def top_counts(counter: Counter[str], limit: int) -> list[dict[str, Any]]:
    return [{"value": value, "count": count} for value, count in counter.most_common(limit)]


def summarize_route(records: list[Trace], route: str, offsets: list[int], top: int) -> dict[str, Any]:
    matched = []
    for record in records:
        if route_for_record(record) == route:
            matched.append(record)

    offset_rows = []
    for offset in offsets:
        buckets: Counter[str] = Counter()
        errors = []
        examples = []
        for record in matched:
            window = record.get("window") or []
            first_right = first_right_approach_index(window)
            if first_right is None:
                row = None
            else:
                index = first_right + offset
                row = window[index] if 0 <= index < len(window) else None
            buckets[row_bucket(row)] += 1
            if row is not None:
                errors.append(float(row["self_paddle_y"]) - float(row["target_y"]))
                if len(examples) < top:
                    examples.append(
                        {
                            "seed": record["seed"],
                            "point": record["point_index"],
                            "outcome": record["outcome"],
                            "pattern": record["pattern"],
                            "row": compact_row(row),
                        }
                    )
        offset_rows.append(
            {
                "offset_from_first_right": offset,
                "top_buckets": top_counts(buckets, top),
                "paddle_error": numeric_summary(errors),
                "examples": examples,
            }
        )

    return {
        "route": route,
        "count": len(matched),
        "outcomes": dict(Counter(record["outcome"] for record in matched)),
        "patterns": top_counts(Counter(record["pattern"] for record in matched), top),
        "seeds": sorted({record["seed"] for record in matched}),
        "offsets": offset_rows,
    }


def route_overview(records: list[Trace], top: int) -> list[dict[str, Any]]:
    by_route: dict[str, list[Trace]] = defaultdict(list)
    for record in records:
        route = route_for_record(record)
        if route is not None:
            by_route[route].append(record)
    rows = []
    for route, items in by_route.items():
        rows.append(
            {
                "route": route,
                "count": len(items),
                "outcomes": dict(Counter(record["outcome"] for record in items)),
                "patterns": top_counts(Counter(record["pattern"] for record in items), top),
                "seeds": sorted({record["seed"] for record in items}),
            }
        )
    rows.sort(key=lambda row: row["count"], reverse=True)
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--route", action="append", default=[])
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--offsets", default="-24,-20,-16,-12,-8,-4,0,4,8,12,16,20,24")
    parser.add_argument("--top", type=int, default=8)
    args = parser.parse_args()

    offsets = [int(part.strip()) for part in args.offsets.split(",") if part.strip()]
    records = read_jsonl(args.jsonl)
    payload = {
        "source": str(args.jsonl),
        "offsets": offsets,
        "top_routes": route_overview(records, args.top)[: args.top],
        "routes": [summarize_route(records, route, offsets, args.top) for route in args.route],
        "interpretation": (
            "Offsets are sampled relative to the first active right-moving frame with ball_x >= 120. "
            "Use matching route summaries to compare whether loss and win routes diverge before "
            "late contact-zone actions."
        ),
    }
    write_json(args.output, payload)
    for route in payload["routes"]:
        patterns = ";".join(f"{item['value']}:{item['count']}" for item in route["patterns"])
        print(f"{route['count']},{route['outcomes']},{route['seeds']},{patterns},{route['route']}")
    print(f"wrote: {args.output}")


if __name__ == "__main__":
    main()
