"""Analyze point-level Pong traces for collision and paddle phase clues."""

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 right_approach(window: list[Trace]) -> list[Trace]:
    return [
        row
        for row in window
        if row.get("dx") is not None and row["dx"] > 0 and row["ball_x"] >= 120
    ]


def contact_zone_sample(window: list[Trace]) -> Trace | None:
    zone = [row for row in right_approach(window) if 184 <= row["ball_x"] <= 195]
    if not zone:
        return None
    return zone[-1]


def paddle_error(row: Trace) -> float:
    return float(row["self_paddle_y"]) - float(row["target_y"])


def summary(values: list[float]) -> dict[str, float | int | None]:
    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 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 analyze_points(records: list[Trace], top: int) -> dict[str, Any]:
    by_outcome = defaultdict(list)
    by_seed = defaultdict(list)
    for record in records:
        by_outcome[record["outcome"]].append(record)
        by_seed[record["seed"]].append(record)

    outcome_summary: dict[str, Any] = {}
    for outcome, rows in sorted(by_outcome.items()):
        final_x = [row["final_ball_x"] for row in rows if row.get("final_ball_x") is not None]
        final_y = [row["final_ball_y"] for row in rows if row.get("final_ball_y") is not None]
        patterns = Counter(row["pattern"] for row in rows)
        zone_errors = []
        zone_x = Counter()
        families: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
        for row in rows:
            sample = contact_zone_sample(row["window"])
            if sample is None:
                continue
            err = paddle_error(sample)
            zone_errors.append(err)
            zone_x[sample["ball_x"]] += 1
            family = (
                sample.get("dy"),
                round(float(sample["target_y"])),
                sample["ball_x"],
            )
            families[family].append(
                {
                    "seed": row["seed"],
                    "point": row["point_index"],
                    "error": err,
                    "pattern": row["pattern"],
                    "reward": row["reward"],
                }
            )

        outcome_summary[outcome] = {
            "count": len(rows),
            "final_x": summary([float(value) for value in final_x]),
            "final_y": summary([float(value) for value in final_y]),
            "final_x_counts": top_counts(Counter(final_x), top),
            "patterns": top_counts(patterns, top),
            "contact_zone_x_counts": top_counts(zone_x, top),
            "contact_zone_error": summary(zone_errors),
            "contact_zone_families": family_summary(families, top),
        }

    seed_summary = {}
    for seed, rows in sorted(by_seed.items()):
        losses = [row for row in rows if row["outcome"] == "loss"]
        seed_summary[str(seed)] = {
            "points": len(rows),
            "score": sum(int(row["reward"]) for row in rows),
            "wins": sum(1 for row in rows if row["outcome"] == "win"),
            "losses": len(losses),
            "loss_patterns": top_counts(Counter(row["pattern"] for row in losses), top),
        }

    return {
        "records": len(records),
        "outcomes": outcome_summary,
        "seeds": seed_summary,
        "blind_spot_signal": blind_spot_signal(outcome_summary),
    }


def family_summary(families: dict[tuple[Any, ...], list[dict[str, Any]]], top: int) -> list[dict[str, Any]]:
    rows = []
    for family, items in families.items():
        rows.append(
            {
                "dy": family[0],
                "target_y_rounded": family[1],
                "ball_x": family[2],
                "count": len(items),
                "mean_error": mean(item["error"] for item in items),
                "examples": items[:5],
            }
        )
    rows.sort(key=lambda row: row["count"], reverse=True)
    return rows[:top]


def blind_spot_signal(outcome_summary: dict[str, Any]) -> dict[str, Any]:
    wins = outcome_summary.get("win", {})
    losses = outcome_summary.get("loss", {})
    return {
        "wins_end_before_policy_paddle_x": wins.get("final_x_counts", []),
        "losses_mostly_past_policy_paddle_x": losses.get("final_x_counts", []),
        "interpretation": (
            "Wins usually register the right-paddle interaction around RAM x 186..190, "
            "while most losses are still right-moving at x 202..204. The policy's "
            "SELF_PADDLE_X_RAM=192 is not by itself a reliable collision boundary; "
            "late target offsets also miss action/phase effects."
        ),
    }


def analyze_action_effects(records: list[Trace], top: int) -> dict[str, Any]:
    by_action: dict[str, list[int]] = defaultdict(list)
    by_action_phase: dict[tuple[str, int], list[int]] = defaultdict(list)
    for record in records:
        window = record["window"]
        for before, after in zip(window, window[1:]):
            if before["self_score"] != after["self_score"]:
                continue
            if before["opponent_score"] != after["opponent_score"]:
                continue
            delta = int(after["self_paddle_y"]) - int(before["self_paddle_y"])
            action = str(before["action_name"])
            phase = int(before["self_paddle_y"]) % 4
            by_action[action].append(delta)
            by_action_phase[(action, phase)].append(delta)

    action_summary = {}
    for action, deltas in sorted(by_action.items()):
        action_summary[action] = {
            "delta": summary([float(value) for value in deltas]),
            "top_deltas": top_counts(Counter(deltas), top),
        }

    phase_summary = []
    for (action, phase), deltas in sorted(by_action_phase.items()):
        if len(deltas) < 30:
            continue
        phase_summary.append(
            {
                "action": action,
                "paddle_y_mod_4": phase,
                "count": len(deltas),
                "top_deltas": top_counts(Counter(deltas), top),
            }
        )

    return {
        "by_action": action_summary,
        "by_action_and_paddle_phase": phase_summary,
        "interpretation": (
            "The same action maps to several next-frame paddle deltas, and NOOP often "
            "continues moving the paddle. This supports treating action effects as "
            "stateful/phase-dependent rather than instantaneous target following."
        ),
    }


def print_report(payload: dict[str, Any]) -> None:
    print(f"records: {payload['points']['records']}")
    print("seed,points,wins,losses,score,loss_patterns")
    for seed, row in payload["points"]["seeds"].items():
        patterns = ";".join(
            f"{item['value']}:{item['count']}" for item in row["loss_patterns"]
        )
        print(f"{seed},{row['points']},{row['wins']},{row['losses']},{row['score']},{patterns}")

    print()
    print("outcome,count,final_x_mean,final_x_min,final_x_max,contact_zone_error_mean")
    for outcome, row in payload["points"]["outcomes"].items():
        final_x = row["final_x"]
        contact_error = row["contact_zone_error"]
        print(
            f"{outcome},{row['count']},{final_x['mean']},"
            f"{final_x['min']},{final_x['max']},{contact_error['mean']}"
        )

    print()
    print("top action effects")
    print("action,count,mean_delta,top_deltas")
    for action, row in payload["action_effects"]["by_action"].items():
        top = ";".join(f"{item['value']}:{item['count']}" for item in row["top_deltas"])
        print(f"{action},{row['delta']['count']},{row['delta']['mean']},{top}")

    print()
    print(payload["points"]["blind_spot_signal"]["interpretation"])
    print(payload["action_effects"]["interpretation"])


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze point traces for phase and collision signals.")
    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/point_phase_summary.json"))
    parser.add_argument("--top", type=int, default=10)
    args = parser.parse_args()

    records = read_jsonl(args.jsonl)
    payload = {
        "source": str(args.jsonl),
        "points": analyze_points(records, top=args.top),
        "action_effects": analyze_action_effects(records, top=args.top),
    }
    write_json(args.output, payload)
    print_report(payload)
    print()
    print(f"wrote: {args.output}")


if __name__ == "__main__":
    main()
