"""Group point terminal states by the next point's reset/initial phase."""

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 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 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']), 2)};"
        f"err={round(float(record['final_paddle_error']), 2)};"
        f"outcome={record['outcome']};pattern={record['pattern']}"
    )


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


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


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 compact_initial(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"),
        "paddle_y": row.get("self_paddle_y"),
        "target_y": row.get("target_y"),
        "action": row.get("action_name"),
        "score": f"{row.get('self_score')}-{row.get('opponent_score')}",
    }


def summarize_links(links: list[tuple[Trace, Trace, Trace | None]], top: int) -> dict[str, Any]:
    next_initial_pads = [
        float(initial["self_paddle_y"])
        for _current, _next_record, initial in links
        if initial is not None
    ]
    next_routes = Counter(route_for_record(next_record) or "none" for _current, next_record, _initial in links)
    next_scores = Counter(score_before(next_record) for _current, next_record, _initial in links)
    current_scores = Counter(score_after(current) for current, _next_record, _initial in links)
    examples = []
    for current, next_record, initial in links[:top]:
        examples.append(
            {
                "seed": current["seed"],
                "point": current["point_index"],
                "score_before": score_before(current),
                "score_after": score_after(current),
                "next_point": next_record["point_index"],
                "next_initial": compact_initial(initial),
                "next_route": route_for_record(next_record),
                "next_outcome": next_record["outcome"],
                "next_pattern": next_record["pattern"],
            }
        )
    return {
        "count": len(links),
        "seeds": sorted({current["seed"] for current, _next_record, _initial in links}),
        "current_scores_after": [
            {"value": value, "count": count}
            for value, count in current_scores.most_common(top)
        ],
        "next_scores_before": [
            {"value": value, "count": count}
            for value, count in next_scores.most_common(top)
        ],
        "next_initial_paddle_y": numeric(next_initial_pads),
        "next_routes": [
            {"value": value, "count": count}
            for value, count in next_routes.most_common(top)
        ],
        "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", default="")
    parser.add_argument("--top", type=int, default=20)
    args = parser.parse_args()

    rows = read_jsonl(args.jsonl)
    by_seed_point = {(record["seed"], record["point_index"]): record for record in rows}
    links_by_terminal: dict[str, list[tuple[Trace, Trace, Trace | None]]] = defaultdict(list)
    for record in rows:
        next_record = by_seed_point.get((record["seed"], record["point_index"] + 1))
        if next_record is None:
            continue
        initial = first_initial_row(next_record)
        links_by_terminal[terminal_key(record)].append((record, next_record, initial))

    summaries = {
        key: summarize_links(links, args.top)
        for key, links in links_by_terminal.items()
    }
    top_terminals = sorted(
        (
            {"terminal_key": key, **summary}
            for key, summary in summaries.items()
        ),
        key=lambda item: item["count"],
        reverse=True,
    )[: args.top]

    selected = None
    if args.terminal_key:
        selected = summaries.get(args.terminal_key)

    payload = {
        "source": str(args.jsonl),
        "terminal_count": len(summaries),
        "top_terminals": top_terminals,
        "selected_terminal_key": args.terminal_key or None,
        "selected": selected,
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    if args.terminal_key:
        print(json.dumps({"selected": selected}, indent=2, sort_keys=True))
    else:
        for item in top_terminals[:10]:
            print(
                f"{item['count']} next_pad={item['next_initial_paddle_y']} "
                f"{item['terminal_key']}"
            )


if __name__ == "__main__":
    main()
