"""Analyze long no-reward action loops in Pong RAM rollouts."""

from __future__ import annotations

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

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 LoopRow:
    step: int
    ball_x: int
    ball_y: int
    self_paddle_y: int
    dx: int | None
    dy: int | None
    target_y: float
    action_name: str
    mode: str
    self_score: int
    opponent_score: int
    score: float


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 signature(row: LoopRow) -> tuple[Any, ...]:
    return (
        row.ball_x,
        row.ball_y,
        row.self_paddle_y,
        row.dx,
        row.dy,
        row.target_y,
        row.action_name,
        row.mode,
        row.self_score,
        row.opponent_score,
    )


def run_seed(seed: int, max_steps: int, tail: int, count_after_step: int | None) -> dict[str, Any]:
    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
        score = 0.0
        last_reward_step = 0
        last_reward_score = 0.0
        rows_tail: deque[LoopRow] = deque(maxlen=tail)
        after_reward_actions: Counter[str] = Counter()
        after_reward_modes: Counter[str] = Counter()
        after_reward_signatures: Counter[tuple[Any, ...]] = Counter()
        inactive_steps_after_last_reward = 0
        done_step: int | None = None

        rows_since_reward: list[LoopRow] = []
        for step in range(1, max_steps + 1):
            state = detect_state(info)
            action = policy.select_action(info)
            dx, dy = dx_dy(state, previous_state)
            row = LoopRow(
                step=step,
                ball_x=state.ball_x,
                ball_y=state.ball_y,
                self_paddle_y=state.self_paddle_y,
                dx=dx,
                dy=dy,
                target_y=float(policy.last_target_y),
                action_name=ACTION_NAMES.get(action, str(action)),
                mode=policy.last_mode,
                self_score=state.self_score,
                opponent_score=state.opponent_score,
                score=score,
            )
            rows_tail.append(row)
            rows_since_reward.append(row)
            previous_state = state

            _obs, reward, done, info = step_env(env, action)
            if reward != 0:
                score += reward
                last_reward_step = step
                last_reward_score = score
                rows_since_reward = []
            if done:
                done_step = step
                break

        analysis_start = count_after_step if count_after_step is not None else last_reward_step
        for row in rows_since_reward:
            if row.step <= analysis_start:
                continue
            after_reward_actions[row.action_name] += 1
            after_reward_modes[row.mode] += 1
            after_reward_signatures[signature(row)] += 1
            if row.mode == "inactive":
                inactive_steps_after_last_reward += 1

        return {
            "seed": seed,
            "max_steps": max_steps,
            "done_step": done_step,
            "final_score": score,
            "last_reward_step": last_reward_step,
            "last_reward_score": last_reward_score,
            "analyzed_after_step": analysis_start,
            "steps_after_last_reward": 0 if done_step is None else done_step - last_reward_step,
            "inactive_steps_after_analysis_start": inactive_steps_after_last_reward,
            "action_counts_after_analysis_start": dict(after_reward_actions.most_common()),
            "mode_counts_after_analysis_start": dict(after_reward_modes.most_common()),
            "top_signatures_after_analysis_start": [
                {"count": count, "signature": list(sig)}
                for sig, count in after_reward_signatures.most_common(12)
            ],
            "tail": [asdict(row) for row in rows_tail],
        }
    finally:
        env.close()


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


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze long no-reward action loops.")
    parser.add_argument("--seeds", default="0")
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument("--tail", type=int, default=80)
    parser.add_argument("--count-after-step", type=int, default=None)
    parser.add_argument("--output", type=Path, default=Path("diagnostics/action_loop_analysis.json"))
    args = parser.parse_args()

    records = [
        run_seed(seed, args.max_steps, args.tail, args.count_after_step)
        for seed in parse_seeds(args.seeds)
    ]
    payload = {
        "description": "Long no-reward action loop analysis for Pong RAM policy rollouts.",
        "records": records,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    print("seed,final_score,done_step,last_reward_step,steps_after_last_reward,top_modes,top_actions")
    for record in records:
        print(
            f"{record['seed']},{record['final_score']},{record['done_step']},"
            f"{record['last_reward_step']},{record['steps_after_last_reward']},"
            f"{record['mode_counts_after_analysis_start']},"
            f"{record['action_counts_after_analysis_start']}"
        )
    print(f"wrote {args.output}")


if __name__ == "__main__":
    main()
