"""Trace Pong RAM policy decisions grouped by point outcome."""

from __future__ import annotations

import argparse
import json
from collections import Counter, deque
from dataclasses import asdict, dataclass
from pathlib import Path
import sys

import envpool
import numpy as np


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from pong_policy import ACTION_NAMES, RamPolicy
from pong_state import PongState, detect_state


@dataclass(frozen=True)
class TraceRow:
    step: int
    ball_x: int
    ball_y: int
    self_paddle_y: int
    opponent_paddle_y: int
    dx: int | None
    dy: int | None
    target_y: float
    action: int
    action_name: str
    mode: str
    low_serve_steps: int
    self_score: int
    opponent_score: int


@dataclass(frozen=True)
class PointRecord:
    seed: int
    point_index: int
    step: int
    outcome: str
    reward: int
    self_score: int
    opponent_score: int
    score_diff: int
    pattern: str
    final_ball_x: int | None
    final_ball_y: int | None
    final_dx: int | None
    final_dy: int | None
    final_target_y: float | None
    final_self_paddle_y: int | None
    final_paddle_error: float | None
    window: list[dict]


def reset_env(env: object) -> tuple[np.ndarray, dict]:
    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        obs, info = reset_out
    else:
        obs, info = reset_out, {}
    return np.asarray(obs), info


def step_env(env: object, action: int) -> tuple[np.ndarray, float, bool, dict]:
    step_out = env.step(np.array([action], dtype=np.int32))
    if len(step_out) == 5:
        obs, reward, terminated, truncated, info = step_out
        done = np.logical_or(terminated, truncated)
    else:
        obs, reward, done, info = step_out
    return (
        np.asarray(obs),
        float(np.asarray(reward).reshape(-1)[0]),
        bool(np.asarray(done).reshape(-1)[0]),
        info,
    )


def dx_dy(state: PongState, previous_state: PongState | None) -> tuple[int | None, int | None]:
    if previous_state is None:
        return None, None
    return state.ball_x - previous_state.ball_x, state.ball_y - previous_state.ball_y


def classify_loss(window: list[TraceRow]) -> str:
    approach = [row for row in window if row.dx is not None and row.dx > 0 and row.ball_x >= 120]
    if not approach:
        return "loss/no_right_approach"

    last = approach[-1]
    paddle_error = last.self_paddle_y - last.target_y
    if last.dy is not None and last.dy > 0 and last.ball_y >= 165:
        vertical = "high_down"
    elif last.dy is not None and last.dy < 0 and last.ball_y <= 120:
        vertical = "low_up"
    elif last.dy is not None and last.dy > 0:
        vertical = "mid_down"
    elif last.dy is not None and last.dy < 0:
        vertical = "mid_up"
    else:
        vertical = "unknown"

    if paddle_error > 8:
        miss = "paddle_below_target"
    elif paddle_error < -8:
        miss = "paddle_above_target"
    else:
        miss = "near_target"
    return f"loss/{vertical}/{miss}"


def point_record(
    seed: int,
    point_index: int,
    step: int,
    reward: int,
    state_after_point: PongState,
    window: list[TraceRow],
) -> PointRecord:
    outcome = "win" if reward > 0 else "loss"
    approach = [row for row in window if row.dx is not None and row.dx > 0 and row.ball_x >= 120]
    final = approach[-1] if approach else (window[-1] if window else None)
    final_paddle_error = None
    if final is not None:
        final_paddle_error = final.self_paddle_y - final.target_y
    return PointRecord(
        seed=seed,
        point_index=point_index,
        step=step,
        outcome=outcome,
        reward=reward,
        self_score=state_after_point.self_score,
        opponent_score=state_after_point.opponent_score,
        score_diff=state_after_point.self_score - state_after_point.opponent_score,
        pattern=classify_loss(window) if reward < 0 else "win",
        final_ball_x=final.ball_x if final is not None else None,
        final_ball_y=final.ball_y if final is not None else None,
        final_dx=final.dx if final is not None else None,
        final_dy=final.dy if final is not None else None,
        final_target_y=final.target_y if final is not None else None,
        final_self_paddle_y=final.self_paddle_y if final is not None else None,
        final_paddle_error=final_paddle_error,
        window=[asdict(row) for row in window],
    )


def run_seed(seed: int, max_steps: int, window_size: int) -> list[PointRecord]:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        _obs, info = reset_env(env)
        policy = RamPolicy()
        policy.reset()
        previous_state: PongState | None = None
        recent: deque[TraceRow] = deque(maxlen=window_size)
        points: list[PointRecord] = []

        for step in range(1, max_steps + 1):
            state = detect_state(info)
            action = policy.select_action(info)
            target_y = policy.last_target_y
            mode = policy.last_mode
            dx, dy = dx_dy(state, previous_state)
            recent.append(
                TraceRow(
                    step=step,
                    ball_x=state.ball_x,
                    ball_y=state.ball_y,
                    self_paddle_y=state.self_paddle_y,
                    opponent_paddle_y=state.opponent_paddle_y,
                    dx=dx,
                    dy=dy,
                    target_y=target_y,
                    action=action,
                    action_name=ACTION_NAMES.get(action, str(action)),
                    mode=mode,
                    low_serve_steps=policy.low_serve_steps,
                    self_score=state.self_score,
                    opponent_score=state.opponent_score,
                )
            )
            previous_state = state
            _obs, reward, done, info = step_env(env, action)
            state_after_point = detect_state(info)
            if reward != 0:
                points.append(
                    point_record(
                        seed=seed,
                        point_index=len(points) + 1,
                        step=step,
                        reward=int(reward),
                        state_after_point=state_after_point,
                        window=list(recent),
                    )
                )
                recent.clear()
            if done:
                break
        return points
    finally:
        env.close()


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


def write_jsonl(path: Path, records: list[PointRecord]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(asdict(record), separators=(",", ":")))
            handle.write("\n")


def print_summary(records: list[PointRecord], tail: int) -> None:
    by_seed: dict[int, list[PointRecord]] = {}
    for record in records:
        by_seed.setdefault(record.seed, []).append(record)

    print("seed,points,wins,losses,final_score,loss_patterns")
    for seed in sorted(by_seed):
        rows = by_seed[seed]
        wins = sum(1 for row in rows if row.outcome == "win")
        losses = sum(1 for row in rows if row.outcome == "loss")
        final = rows[-1].score_diff if rows else 0
        patterns = Counter(row.pattern for row in rows if row.outcome == "loss")
        pattern_text = ";".join(f"{name}:{count}" for name, count in patterns.most_common())
        print(f"{seed},{len(rows)},{wins},{losses},{final},{pattern_text}")

    print()
    print("loss detail")
    print("seed,point,step,pattern,bx,by,dx,dy,target,paddle,error")
    for record in records:
        if record.outcome != "loss":
            continue
        print(
            f"{record.seed},{record.point_index},{record.step},{record.pattern},"
            f"{record.final_ball_x},{record.final_ball_y},"
            f"{record.final_dx},{record.final_dy},"
            f"{none_safe_float(record.final_target_y)},"
            f"{record.final_self_paddle_y},"
            f"{none_safe_float(record.final_paddle_error)}"
        )

    if tail <= 0:
        return
    print()
    print(f"last {tail} frames before each loss")
    for record in records:
        if record.outcome != "loss":
            continue
        print(f"\nseed={record.seed} point={record.point_index} step={record.step} pattern={record.pattern}")
        print("step,bx,by,pad,dx,dy,target,action,mode")
        for row in record.window[-tail:]:
            print(
                f"{row['step']},{row['ball_x']},{row['ball_y']},{row['self_paddle_y']},"
                f"{row['dx']},{row['dy']},{row['target_y']:.1f},"
                f"{row['action_name']},{row['mode']}"
            )


def none_safe_float(value: float | None) -> str:
    if value is None:
        return ""
    return f"{value:.1f}"


def main() -> None:
    parser = argparse.ArgumentParser(description="Trace Pong policy decisions by point outcome.")
    parser.add_argument("--seeds", default="0,1,3,4,5,8,9")
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument("--window", type=int, default=80)
    parser.add_argument("--tail", type=int, default=40)
    parser.add_argument("--jsonl", type=Path, default=Path("diagnostics/point_traces.jsonl"))
    args = parser.parse_args()
    if args.window < 1:
        raise ValueError("--window must be at least 1")
    if args.tail < 0:
        raise ValueError("--tail must be non-negative")

    records: list[PointRecord] = []
    for seed in parse_seeds(args.seeds):
        records.extend(run_seed(seed=seed, max_steps=args.max_steps, window_size=args.window))

    write_jsonl(args.jsonl, records)
    print(f"envpool: {envpool.__version__}")
    print(f"seeds: {args.seeds}")
    print(f"max_steps: {args.max_steps}")
    print(f"window: {args.window}")
    print(f"wrote: {args.jsonl}")
    print()
    print_summary(records, tail=args.tail)


if __name__ == "__main__":
    main()
