"""Inspect the source window around terminal-reset classifier keys."""

from __future__ import annotations

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


SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from compare_setup_routes import route_for_record


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) + "\n", encoding="utf-8")


def score_before(record: Trace) -> str:
    window = record.get("window") or []
    if not window:
        return "unknown"
    first = window[0]
    return f"{first.get('self_score')}-{first.get('opponent_score')}"


def score_after(record: Trace) -> str:
    return f"{record.get('self_score')}-{record.get('opponent_score')}"


def terminal_key(record: Trace) -> str:
    return (
        f"x={record['final_ball_x']};y={record['final_ball_y']};"
        f"dx={record['final_dx']};dy={record['final_dy']};"
        f"pad={record['final_self_paddle_y']};"
        f"target={round(float(record['final_target_y']), 1)};"
        f"err={round(float(record['final_paddle_error']), 1)};"
        f"outcome={record['outcome']};pattern={record['pattern']}"
    )


def first_initial_row(record: Trace) -> Trace | None:
    for row in record.get("window") or []:
        if row.get("mode") == "initial" and row.get("ball_x") not in (0, 52, 205):
            return row
    return None


def compact_row(row: Trace | None) -> dict[str, Any] | None:
    if row is None:
        return None
    target_y = row.get("target_y")
    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"),
        "paddle_y": row.get("self_paddle_y"),
        "target_y": target_y,
        "paddle_error": None if target_y is None else row.get("self_paddle_y") - target_y,
        "action": row.get("action_name"),
        "mode": row.get("mode"),
        "low_serve_steps": row.get("low_serve_steps"),
        "score": f"{row.get('self_score')}-{row.get('opponent_score')}",
    }


def find_terminal_row_index(record: Trace) -> int | None:
    target = {
        "ball_x": record.get("final_ball_x"),
        "ball_y": record.get("final_ball_y"),
        "dx": record.get("final_dx"),
        "dy": record.get("final_dy"),
        "self_paddle_y": record.get("final_self_paddle_y"),
    }
    for index, row in enumerate(record.get("window") or []):
        if all(row.get(name) == value for name, value in target.items()):
            return index
    return None


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


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


def summarize_offset_rows(rows: list[Trace]) -> dict[str, Any]:
    return {
        "count": len(rows),
        "ball_x": numeric([float(row["ball_x"]) for row in rows if row.get("ball_x") is not None]),
        "ball_y": numeric([float(row["ball_y"]) for row in rows if row.get("ball_y") is not None]),
        "paddle_y": numeric(
            [float(row["self_paddle_y"]) for row in rows if row.get("self_paddle_y") is not None]
        ),
        "target_y": numeric([float(row["target_y"]) for row in rows if row.get("target_y") is not None]),
        "paddle_error": numeric(
            [
                float(row["self_paddle_y"] - row["target_y"])
                for row in rows
                if row.get("self_paddle_y") is not None and row.get("target_y") is not None
            ]
        ),
        "actions": counter_payload(Counter(row.get("action_name") for row in rows)),
        "modes": counter_payload(Counter(row.get("mode") for row in rows)),
        "low_serve_steps": counter_payload(Counter(row.get("low_serve_steps") for row in rows)),
    }


def summarize_matches(
    key: str,
    matches: list[tuple[Trace, int, Trace | None, int | None]],
    radius_before: int,
    radius_after: int,
    top: int,
) -> dict[str, Any]:
    rows_by_offset: dict[int, list[Trace]] = defaultdict(list)
    examples = []
    next_initial_pads = []
    next_routes = Counter()
    next_labels = Counter()
    current_scores = Counter()
    final_scores = Counter()
    for record, terminal_index, next_record, seed_final_score in matches:
        window = record.get("window") or []
        for offset in range(-radius_before, radius_after + 1):
            absolute = terminal_index + offset
            if 0 <= absolute < len(window):
                rows_by_offset[offset].append(window[absolute])

        initial = first_initial_row(next_record) if next_record is not None else None
        if initial is not None:
            next_initial_pads.append(float(initial["self_paddle_y"]))
        if next_record is not None:
            next_routes[route_for_record(next_record) or "none"] += 1
            if next_record.get("outcome") == "loss":
                label = (
                    f"loss:{next_record.get('pattern')};"
                    f"final_y={next_record.get('final_ball_y')};"
                    f"score_after={score_after(next_record)}"
                )
            else:
                label = "win"
            next_labels[label] += 1
        current_scores[score_after(record)] += 1
        final_scores[seed_final_score] += 1

        if len(examples) < top:
            start = max(0, terminal_index - radius_before)
            end = min(len(window), terminal_index + radius_after + 1)
            examples.append(
                {
                    "seed": record["seed"],
                    "point": record["point_index"],
                    "score_before": score_before(record),
                    "score_after": score_after(record),
                    "seed_final_score": seed_final_score,
                    "terminal_index": terminal_index,
                    "terminal_row": compact_row(window[terminal_index]),
                    "next_initial": compact_row(initial),
                    "next_route": None if next_record is None else route_for_record(next_record),
                    "next_outcome": None if next_record is None else next_record.get("outcome"),
                    "next_pattern": None if next_record is None else next_record.get("pattern"),
                    "window": [
                        {"offset": index - terminal_index, **(compact_row(row) or {})}
                        for index, row in enumerate(window[start:end], start=start)
                    ],
                }
            )

    return {
        "terminal_key": key,
        "count": len(matches),
        "seeds": sorted({record["seed"] for record, _index, _next, _final in matches}),
        "current_scores_after": counter_payload(current_scores, top),
        "seed_final_scores": counter_payload(final_scores, top),
        "next_initial_paddle_y": numeric(next_initial_pads),
        "next_routes": counter_payload(next_routes, top),
        "next_labels": counter_payload(next_labels, top),
        "offset_summary": {
            str(offset): summarize_offset_rows(rows)
            for offset, rows in sorted(rows_by_offset.items())
        },
        "examples": examples,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--terminal-key", action="append", default=[])
    parser.add_argument("--before", type=int, default=12)
    parser.add_argument("--after", type=int, default=16)
    parser.add_argument("--top", type=int, default=8)
    args = parser.parse_args()

    records = read_jsonl(args.jsonl)
    by_seed_point = {(record["seed"], record["point_index"]): record for record in records}
    final_score_by_seed = {}
    for seed in sorted({record["seed"] for record in records}):
        seed_records = sorted(
            [record for record in records if record["seed"] == seed],
            key=lambda row: row["point_index"],
        )
        if seed_records:
            final_score_by_seed[seed] = int(seed_records[-1]["score_diff"])

    requested = set(args.terminal_key)
    matches_by_key: dict[str, list[tuple[Trace, int, Trace | None, int | None]]] = defaultdict(list)
    missing_source_rows = []
    for record in records:
        key = terminal_key(record)
        if requested and key not in requested:
            continue
        terminal_index = find_terminal_row_index(record)
        if terminal_index is None:
            missing_source_rows.append({"seed": record["seed"], "point": record["point_index"], "terminal_key": key})
            continue
        next_record = by_seed_point.get((record["seed"], record["point_index"] + 1))
        matches_by_key[key].append(
            (record, terminal_index, next_record, final_score_by_seed.get(record["seed"]))
        )

    payload = {
        "source": str(args.jsonl),
        "before": args.before,
        "after": args.after,
        "requested_terminal_keys": args.terminal_key,
        "missing_source_rows": missing_source_rows[: args.top],
        "groups": [
            summarize_matches(key, matches, args.before, args.after, args.top)
            for key, matches in sorted(matches_by_key.items(), key=lambda item: len(item[1]), reverse=True)
        ],
        "interpretation": (
            "Rows are aligned at the terminal classifier row from the previous point. "
            "Offsets before/after describe the source paddle/action cycle that precedes the next reset phase."
        ),
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    for group in payload["groups"][:5]:
        print(
            f"count={group['count']} seeds={group['seeds']} "
            f"next_pad={group['next_initial_paddle_y']} key={group['terminal_key']}"
        )


if __name__ == "__main__":
    main()
