"""Summarize score/outcome phases within one setup route."""

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_route_outcomes import active_right_indices, compact_row
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_key(record: Trace) -> str:
    window = record.get("window") or []
    if window:
        first = window[0]
        return f"{first.get('self_score')}-{first.get('opponent_score')}"
    return f"{record.get('self_score')}-{record.get('opponent_score')}"


def row_at(record: Trace, anchor: str, offset: int) -> Trace | None:
    window = record.get("window") or []
    indices = active_right_indices(window)
    if not indices:
        return None
    base = indices[0] if anchor == "first_right" else indices[-1]
    index = base + offset
    if 0 <= index < len(window):
        return window[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 row_signature(row: Trace | None) -> str:
    if row is None:
        return "missing"
    return (
        f"x={row.get('ball_x')};y={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 summarize_offsets(records: list[Trace], anchor: str, offsets: list[int]) -> list[dict[str, Any]]:
    rows = []
    for offset in offsets:
        errors: list[float] = []
        signatures: Counter[str] = Counter()
        examples = []
        for record in records:
            row = row_at(record, anchor, offset)
            signatures[row_signature(row)] += 1
            if row is not None:
                errors.append(float(row["self_paddle_y"]) - float(row["target_y"]))
            if row is not None and len(examples) < 3:
                examples.append(
                    {
                        "seed": record["seed"],
                        "point": record["point_index"],
                        "row": compact_row(row),
                    }
                )
        rows.append(
            {
                "offset": offset,
                "paddle_error": numeric(errors),
                "top_signatures": [
                    {"value": value, "count": count}
                    for value, count in signatures.most_common(5)
                ],
                "examples": examples,
            }
        )
    return rows


def summarize_group(
    records: list[Trace],
    first_offsets: list[int],
    last_offsets: list[int],
) -> dict[str, Any]:
    return {
        "count": len(records),
        "seeds": sorted({record["seed"] for record in records}),
        "points": sorted({record["point_index"] for record in records}),
        "patterns": dict(Counter(record["pattern"] for record in records)),
        "final_ball_y": numeric([float(record["final_ball_y"]) for record in records]),
        "final_paddle_error": numeric([float(record["final_paddle_error"]) for record in records]),
        "first_right_offsets": summarize_offsets(records, "first_right", first_offsets),
        "last_right_offsets": summarize_offsets(records, "last_right", last_offsets),
    }


def parse_offsets(raw: str) -> list[int]:
    return [int(part.strip()) for part in raw.split(",") if part.strip()]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--route", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--first-offsets", default="20,24,28,32,36,40,44,48")
    parser.add_argument("--last-offsets", default="-28,-24,-20,-16,-12,-8,-4,0,1,2,3,4")
    args = parser.parse_args()

    first_offsets = parse_offsets(args.first_offsets)
    last_offsets = parse_offsets(args.last_offsets)
    route_records = [
        record for record in read_jsonl(args.jsonl)
        if route_for_record(record) == args.route
    ]

    by_phase: dict[str, list[Trace]] = defaultdict(list)
    for record in route_records:
        by_phase[f"{score_key(record)}:{record['outcome']}:{record['pattern']}"].append(record)

    payload = {
        "source": str(args.jsonl),
        "route": args.route,
        "count": len(route_records),
        "phase_groups": {
            key: summarize_group(records, first_offsets, last_offsets)
            for key, records in sorted(by_phase.items())
        },
    }
    write_json(args.output, payload)
    print(f"wrote: {args.output}")
    for key, summary in payload["phase_groups"].items():
        print(f"{key}: count={summary['count']} seeds={summary['seeds'][:10]}")


if __name__ == "__main__":
    main()
