"""Analyze stateful right-paddle dynamics from 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, Callable


Trace = dict[str, Any]
Transition = dict[str, Any]
FeatureFn = Callable[[Transition], tuple[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 transitions(records: list[Trace]) -> list[Transition]:
    rows: list[Transition] = []
    for record in records:
        window = record["window"]
        previous_delta = None
        for before, after in zip(window, window[1:]):
            if before["self_score"] != after["self_score"]:
                previous_delta = None
                continue
            if before["opponent_score"] != after["opponent_score"]:
                previous_delta = None
                continue
            delta = int(after["self_paddle_y"]) - int(before["self_paddle_y"])
            rows.append(
                {
                    "seed": int(record["seed"]),
                    "outcome": record["outcome"],
                    "pattern": record["pattern"],
                    "point_index": int(record["point_index"]),
                    "step": int(before["step"]),
                    "ball_x": int(before["ball_x"]),
                    "ball_y": int(before["ball_y"]),
                    "dx": before.get("dx"),
                    "dy": before.get("dy"),
                    "target_y": float(before["target_y"]),
                    "paddle_y": int(before["self_paddle_y"]),
                    "paddle_error": float(before["self_paddle_y"]) - float(before["target_y"]),
                    "action": str(before["action_name"]),
                    "mode": str(before["mode"]),
                    "previous_delta": previous_delta,
                    "delta": delta,
                }
            )
            previous_delta = delta
    return rows


def majority_model_accuracy(
    rows: list[Transition],
    feature_fn: FeatureFn,
    train_seeds: set[int],
    test_seeds: set[int],
) -> dict[str, Any]:
    counts: dict[tuple[Any, ...], Counter[int]] = defaultdict(Counter)
    default_counts: Counter[int] = Counter()
    for row in rows:
        if row["seed"] not in train_seeds:
            continue
        counts[feature_fn(row)][int(row["delta"])] += 1
        default_counts[int(row["delta"])] += 1

    default_delta = default_counts.most_common(1)[0][0]
    predictions = {}
    for key, counter in counts.items():
        predictions[key] = counter.most_common(1)[0][0]

    total = 0
    exact = 0
    sign_match = 0
    abs_error = []
    used_default = 0
    for row in rows:
        if row["seed"] not in test_seeds:
            continue
        key = feature_fn(row)
        if key in predictions:
            predicted = predictions[key]
        else:
            predicted = default_delta
            used_default += 1
        actual = int(row["delta"])
        total += 1
        exact += int(predicted == actual)
        sign_match += int(sign(predicted) == sign(actual))
        abs_error.append(abs(predicted - actual))

    return {
        "test_count": total,
        "exact_accuracy": exact / total if total else None,
        "sign_accuracy": sign_match / total if total else None,
        "mean_abs_error": mean(abs_error) if abs_error else None,
        "used_default": used_default,
        "train_feature_count": len(predictions),
    }


def feature_sets() -> dict[str, FeatureFn]:
    return {
        "action": lambda row: (row["action"],),
        "action_paddle_mod4": lambda row: (row["action"], row["paddle_y"] % 4),
        "action_prev_delta_sign": lambda row: (row["action"], sign(row["previous_delta"])),
        "action_paddle_mod4_prev_delta_sign": lambda row: (
            row["action"],
            row["paddle_y"] % 4,
            sign(row["previous_delta"]),
        ),
        "action_paddle_mod4_error_bin_prev_delta_sign": lambda row: (
            row["action"],
            row["paddle_y"] % 4,
            error_bin(float(row["paddle_error"])),
            sign(row["previous_delta"]),
        ),
    }


def error_bin(value: 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 cross_seed_report(rows: list[Transition]) -> dict[str, Any]:
    all_seeds = sorted({int(row["seed"]) for row in rows})
    folds = []
    by_feature = {name: [] for name in feature_sets()}
    for test_seed in all_seeds:
        train_seeds = set(all_seeds) - {test_seed}
        test_seeds = {test_seed}
        fold = {"test_seed": test_seed, "models": {}}
        for name, feature_fn in feature_sets().items():
            metrics = majority_model_accuracy(rows, feature_fn, train_seeds, test_seeds)
            fold["models"][name] = metrics
            by_feature[name].append(metrics)
        folds.append(fold)

    aggregate = {}
    for name, metrics_rows in by_feature.items():
        aggregate[name] = {
            "exact_accuracy": mean(row["exact_accuracy"] for row in metrics_rows if row["exact_accuracy"] is not None),
            "sign_accuracy": mean(row["sign_accuracy"] for row in metrics_rows if row["sign_accuracy"] is not None),
            "mean_abs_error": mean(row["mean_abs_error"] for row in metrics_rows if row["mean_abs_error"] is not None),
            "mean_used_default": mean(float(row["used_default"]) for row in metrics_rows),
            "mean_train_feature_count": mean(float(row["train_feature_count"]) for row in metrics_rows),
        }

    return {"seeds": all_seeds, "folds": folds, "aggregate": aggregate}


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 contact_phase_report(records: list[Trace], top: int) -> dict[str, Any]:
    buckets: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for record in records:
        window = record["window"]
        previous_delta = None
        for before, after in zip(window, window[1:]):
            if before["self_score"] != after["self_score"] or before["opponent_score"] != after["opponent_score"]:
                previous_delta = None
                continue
            delta = int(after["self_paddle_y"]) - int(before["self_paddle_y"])
            if (
                before.get("dx") is not None
                and before["dx"] > 0
                and 184 <= before["ball_x"] <= 195
            ):
                buckets[str(record["outcome"])].append(
                    {
                        "seed": record["seed"],
                        "point": record["point_index"],
                        "pattern": record["pattern"],
                        "ball_x": before["ball_x"],
                        "ball_y": before["ball_y"],
                        "dy": before.get("dy"),
                        "target_y": before["target_y"],
                        "paddle_y": before["self_paddle_y"],
                        "paddle_error": before["self_paddle_y"] - before["target_y"],
                        "action": before["action_name"],
                        "previous_delta": previous_delta,
                        "delta": delta,
                        "phase_key": phase_key(before, previous_delta),
                    }
                )
            previous_delta = delta

    report = {}
    for outcome, rows in sorted(buckets.items()):
        report[outcome] = {
            "count": len(rows),
            "phase_keys": top_counts(Counter(row["phase_key"] for row in rows), top),
            "actions": top_counts(Counter(row["action"] for row in rows), top),
            "previous_delta_signs": top_counts(Counter(sign(row["previous_delta"]) for row in rows), top),
            "next_delta_signs": top_counts(Counter(sign(row["delta"]) for row in rows), top),
            "paddle_error": numeric_summary([float(row["paddle_error"]) for row in rows]),
            "examples": rows[:top],
        }
    return report


def phase_key(row: Trace, previous_delta: int | None) -> str:
    return (
        f"dy={row.get('dy')};"
        f"action={row['action_name']};"
        f"mod4={int(row['self_paddle_y']) % 4};"
        f"prev={sign(previous_delta)};"
        f"err={error_bin(float(row['self_paddle_y']) - float(row['target_y']))}"
    )


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


def print_report(payload: dict[str, Any]) -> None:
    print(f"source: {payload['source']}")
    print(f"transitions: {payload['transition_count']}")
    print("model,exact_accuracy,sign_accuracy,mean_abs_error,mean_feature_count")
    for name, row in payload["cross_seed"]["aggregate"].items():
        print(
            f"{name},{row['exact_accuracy']:.3f},{row['sign_accuracy']:.3f},"
            f"{row['mean_abs_error']:.3f},{row['mean_train_feature_count']:.1f}"
        )
    print()
    print("contact_outcome,count,top_phase_keys")
    for outcome, row in payload["contact_phase"].items():
        top = "|".join(f"{item['value']}:{item['count']}" for item in row["phase_keys"][:5])
        print(f"{outcome},{row['count']},{top}")
    print()
    print(payload["interpretation"])


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze stateful paddle dynamics in point traces.")
    parser.add_argument(
        "--jsonl",
        type=Path,
        default=Path("diagnostics/point_traces_seeds_0_1_3_4_5_8_9.jsonl"),
    )
    parser.add_argument("--output", type=Path, default=Path("diagnostics/paddle_dynamics_summary.json"))
    parser.add_argument("--top", type=int, default=10)
    args = parser.parse_args()

    records = read_jsonl(args.jsonl)
    rows = transitions(records)
    payload = {
        "source": str(args.jsonl),
        "transition_count": len(rows),
        "cross_seed": cross_seed_report(rows),
        "contact_phase": contact_phase_report(records, top=args.top),
        "interpretation": (
            "Cross-seed majority prediction estimates whether paddle motion is better "
            "explained by action alone or by action plus phase/history. Contact-zone "
            "buckets expose whether wins and losses occupy different phase keys before "
            "the reward frame."
        ),
    }
    write_json(args.output, payload)
    print_report(payload)
    print(f"wrote: {args.output}")


if __name__ == "__main__":
    main()
