"""Compare paddle/action cycles around an anchor 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_setup_routes import compact_row, 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 phase_key(record: Trace) -> str:
    return f"{score_key(record)}:{record['outcome']}:{record['pattern']}"


def find_anchor_index(
    window: list[Trace],
    anchor_x: int,
    anchor_y: int,
    anchor_dx: int | None,
    anchor_dy: int | None,
) -> int | None:
    for index, row in enumerate(window):
        if row.get("ball_x") != anchor_x or row.get("ball_y") != anchor_y:
            continue
        if anchor_dx is not None and row.get("dx") != anchor_dx:
            continue
        if anchor_dy is not None and row.get("dy") != anchor_dy:
            continue
        return index
    return None


def row_at(record: Trace, anchor_index: int, offset: int) -> Trace | None:
    window = record.get("window") or []
    index = anchor_index + offset
    if 0 <= index < len(window):
        return window[index]
    return None


def previous_delta(record: Trace, anchor_index: int, offset: int) -> int | None:
    window = record.get("window") or []
    index = anchor_index + offset
    if index <= 0 or index >= len(window):
        return None
    return int(window[index]["self_paddle_y"]) - int(window[index - 1]["self_paddle_y"])


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, delta: int | None) -> str:
    if row is None:
        return "missing"
    err = round(float(row["self_paddle_y"]) - float(row["target_y"]), 2)
    return (
        f"x={row.get('ball_x')};y={row.get('ball_y')};"
        f"dx={row.get('dx')};dy={row.get('dy')};"
        f"pad={row.get('self_paddle_y')};delta={delta};"
        f"target={round(float(row['target_y']), 2)};err={err};"
        f"action={row.get('action_name')};mode={row.get('mode')}"
    )


def sequence_signature(record: Trace, anchor_index: int, offsets: list[int]) -> str:
    pieces = []
    for offset in offsets:
        row = row_at(record, anchor_index, offset)
        if row is None:
            pieces.append(f"{offset}:missing")
        else:
            pieces.append(
                f"{offset}:x{row.get('ball_x')}/y{row.get('ball_y')}/"
                f"p{row.get('self_paddle_y')}/d{previous_delta(record, anchor_index, offset)}/"
                f"{row.get('action_name')}/{row.get('mode')}"
            )
    return "|".join(pieces)


def summarize_group(
    records: list[tuple[Trace, int]],
    offsets: list[int],
    sequence_offsets: list[int],
    top: int,
) -> dict[str, Any]:
    offset_rows = []
    for offset in offsets:
        signatures: Counter[str] = Counter()
        paddles: list[float] = []
        deltas: list[float] = []
        errors: list[float] = []
        examples = []
        for record, anchor_index in records:
            row = row_at(record, anchor_index, offset)
            delta = previous_delta(record, anchor_index, offset)
            signatures[row_signature(row, delta)] += 1
            if row is not None:
                paddles.append(float(row["self_paddle_y"]))
                errors.append(float(row["self_paddle_y"]) - float(row["target_y"]))
                if delta is not None:
                    deltas.append(float(delta))
                if len(examples) < top:
                    examples.append(
                        {
                            "seed": record["seed"],
                            "point": record["point_index"],
                            "row": compact_row(row),
                            "previous_paddle_delta": delta,
                        }
                    )
        offset_rows.append(
            {
                "offset": offset,
                "paddle_y": numeric(paddles),
                "previous_paddle_delta": numeric(deltas),
                "paddle_error": numeric(errors),
                "top_signatures": [
                    {"value": value, "count": count}
                    for value, count in signatures.most_common(top)
                ],
                "examples": examples,
            }
        )

    sequence_counts = Counter(
        sequence_signature(record, anchor_index, sequence_offsets)
        for record, anchor_index in records
    )
    return {
        "count": len(records),
        "seeds": sorted({record["seed"] for record, _anchor_index in records}),
        "points": sorted({record["point_index"] for record, _anchor_index in records}),
        "offsets": offset_rows,
        "top_cycle_signatures": [
            {"value": value, "count": count}
            for value, count in sequence_counts.most_common(top)
        ],
    }


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


def parse_optional_int(raw: str) -> int | None:
    if raw.lower() in {"none", "any", ""}:
        return None
    return int(raw)


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("--anchor-x", type=int, required=True)
    parser.add_argument("--anchor-y", type=int, required=True)
    parser.add_argument("--anchor-dx", default="any")
    parser.add_argument("--anchor-dy", default="any")
    parser.add_argument("--offsets", default="-32,-28,-24,-20,-16,-12,-8,-4,0,4,8")
    parser.add_argument("--sequence-offsets", default="-24,-20,-16,-12,-8,-4,0")
    parser.add_argument("--top", type=int, default=8)
    args = parser.parse_args()

    offsets = parse_offsets(args.offsets)
    sequence_offsets = parse_offsets(args.sequence_offsets)
    anchor_dx = parse_optional_int(args.anchor_dx)
    anchor_dy = parse_optional_int(args.anchor_dy)

    by_phase: dict[str, list[tuple[Trace, int]]] = defaultdict(list)
    missing_anchor: Counter[str] = Counter()
    matched_route_count = 0
    for record in read_jsonl(args.jsonl):
        if route_for_record(record) != args.route:
            continue
        matched_route_count += 1
        window = record.get("window") or []
        anchor_index = find_anchor_index(window, args.anchor_x, args.anchor_y, anchor_dx, anchor_dy)
        key = phase_key(record)
        if anchor_index is None:
            missing_anchor[key] += 1
            continue
        by_phase[key].append((record, anchor_index))

    payload = {
        "source": str(args.jsonl),
        "route": args.route,
        "anchor": {
            "x": args.anchor_x,
            "y": args.anchor_y,
            "dx": anchor_dx,
            "dy": anchor_dy,
        },
        "matched_route_count": matched_route_count,
        "missing_anchor_by_phase": dict(missing_anchor),
        "phase_groups": {
            key: summarize_group(records, offsets, sequence_offsets, args.top)
            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'][:12]}")


if __name__ == "__main__":
    main()
