"""Analyze pre-return setup phase keys in point trace windows."""

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 sign(value: int | float | None) -> int:
    if value is None:
        return 0
    if value < 0:
        return -1
    if value > 0:
        return 1
    return 0


def paddle_delta(window: list[Trace], index: int) -> int | None:
    if index <= 0:
        return None
    return int(window[index]["self_paddle_y"]) - int(window[index - 1]["self_paddle_y"])


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:
            return index
    return None


def setup_samples(record: Trace) -> dict[str, Any] | None:
    window = record["window"]
    first_right = first_right_approach_index(window)
    if first_right is None:
        return None

    before_indices = [index for index in range(max(0, first_right - 24), first_right)]
    right_indices = [index for index in range(first_right, min(len(window), first_right + 12))]
    if not right_indices:
        return None

    entry = window[first_right]
    pre = window[before_indices[-1]] if before_indices else entry
    right = window[right_indices[min(3, len(right_indices) - 1)]]

    return {
        "seed": record["seed"],
        "point": record["point_index"],
        "outcome": record["outcome"],
        "pattern": record["pattern"],
        "reward": record["reward"],
        "score_before": f"{entry['self_score']}-{entry['opponent_score']}",
        "entry": sample_row(window, first_right),
        "pre_entry": sample_row(window, before_indices[-1]) if before_indices else sample_row(window, first_right),
        "right_early": sample_row(window, right_indices[min(3, len(right_indices) - 1)]),
        "keys": {
            "entry": phase_key(entry, paddle_delta(window, first_right)),
            "pre_entry": phase_key(pre, paddle_delta(window, before_indices[-1])) if before_indices else phase_key(entry, None),
            "right_early": phase_key(right, paddle_delta(window, right_indices[min(3, len(right_indices) - 1)])),
            "route": route_key(pre, entry, right),
        },
    }


def sample_row(window: list[Trace], index: int) -> dict[str, Any]:
    row = window[index]
    return {
        "step": row["step"],
        "ball_x": row["ball_x"],
        "ball_y": row["ball_y"],
        "self_paddle_y": row["self_paddle_y"],
        "dx": row.get("dx"),
        "dy": row.get("dy"),
        "target_y": row["target_y"],
        "action": row["action_name"],
        "mode": row["mode"],
        "paddle_delta": paddle_delta(window, index),
        "error": row["self_paddle_y"] - row["target_y"],
    }


def phase_key(row: Trace, delta: int | None) -> str:
    return (
        f"bx={bucket(row['ball_x'], 10)};"
        f"by={bucket(row['ball_y'], 10)};"
        f"pad={bucket(row['self_paddle_y'], 10)};"
        f"dx={row.get('dx')};dy={row.get('dy')};"
        f"a={row['action_name']};"
        f"pd={sign(delta)};"
        f"err={error_bucket(row['self_paddle_y'] - row['target_y'])}"
    )


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


def error_bucket(value: int | float) -> str:
    if value < -16:
        return "<-16"
    if value < -8:
        return "-16..-9"
    if value <= 8:
        return "-8..8"
    if value <= 16:
        return "9..16"
    return ">16"


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


def summarize(samples: list[dict[str, Any]], top: int) -> dict[str, Any]:
    groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for sample in samples:
        groups[f"{sample['outcome']}:{sample['pattern']}"].append(sample)

    group_rows = {}
    for name, rows in sorted(groups.items()):
        group_rows[name] = {
            "count": len(rows),
            "seeds": sorted({row["seed"] for row in rows}),
            "entry_keys": top_counts(Counter(row["keys"]["entry"] for row in rows), top),
            "pre_entry_keys": top_counts(Counter(row["keys"]["pre_entry"] for row in rows), top),
            "right_early_keys": top_counts(Counter(row["keys"]["right_early"] for row in rows), top),
            "route_keys": top_counts(Counter(row["keys"]["route"] for row in rows), top),
            "entry_error_mean": mean(row["entry"]["error"] for row in rows),
            "examples": rows[: min(top, len(rows))],
        }

    by_route = defaultdict(list)
    for sample in samples:
        by_route[sample["keys"]["route"]].append(sample)
    mixed_routes = []
    for route, rows in by_route.items():
        outcomes = Counter(row["outcome"] for row in rows)
        patterns = Counter(row["pattern"] for row in rows)
        if len(rows) >= 3:
            mixed_routes.append(
                {
                    "route": route,
                    "count": len(rows),
                    "outcomes": dict(outcomes),
                    "patterns": top_counts(patterns, top),
                    "seeds": sorted({row["seed"] for row in rows}),
                }
            )
    mixed_routes.sort(key=lambda row: row["count"], reverse=True)

    return {
        "sample_count": len(samples),
        "groups": group_rows,
        "routes": mixed_routes[:top],
    }


def print_report(payload: dict[str, Any]) -> None:
    print(f"source: {payload['source']}")
    print(f"samples: {payload['summary']['sample_count']}")
    print("group,count,seeds,top_route")
    for name, row in payload["summary"]["groups"].items():
        route = row["route_keys"][0]["value"] if row["route_keys"] else ""
        print(f"{name},{row['count']},{row['seeds']},{route}")
    print()
    print("top routes")
    for row in payload["summary"]["routes"][:10]:
        print(row)


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze pre-return setup phase from point traces.")
    parser.add_argument("--jsonl", type=Path, required=True)
    parser.add_argument("--output", type=Path, default=Path("diagnostics/setup_phase_summary.json"))
    parser.add_argument("--top", type=int, default=8)
    args = parser.parse_args()

    records = read_jsonl(args.jsonl)
    samples = [sample for record in records if (sample := setup_samples(record)) is not None]
    payload = {
        "source": str(args.jsonl),
        "summary": summarize(samples, top=args.top),
        "interpretation": (
            "Setup phase keys are sampled around the first right-moving return in each point. "
            "Use this to compare whether repeated losses are already distinguishable before "
            "late right-edge contact."
        ),
    }
    write_json(args.output, payload)
    print_report(payload)
    print(f"wrote: {args.output}")


if __name__ == "__main__":
    main()
