"""Trace RAM policy decisions around score events."""

from __future__ import annotations

import argparse
from collections import deque
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


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 main() -> None:
    parser = argparse.ArgumentParser(description="Trace policy decisions around Pong score events.")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--max-steps", type=int, default=10_000)
    parser.add_argument("--window", type=int, default=24)
    args = parser.parse_args()

    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=args.seed)
    try:
        _obs, info = reset_env(env)
        policy = RamPolicy()
        policy.reset()
        previous_state = None
        score = 0.0
        recent = deque(maxlen=args.window)

        for step in range(1, args.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)
            _obs, reward, done, info = step_env(env, action)
            score += reward
            recent.append(
                {
                    "step": step,
                    "ball_x": state.ball_x,
                    "ball_y": state.ball_y,
                    "paddle_y": state.self_paddle_y,
                    "target_y": target_y,
                    "dx": dx,
                    "dy": dy,
                    "action": action,
                    "mode": mode,
                    "low_serve_steps": policy.low_serve_steps,
                    "reward": reward,
                    "score": score,
                }
            )
            previous_state = state

            if reward != 0:
                print(f"\nreward event step={step} reward={reward:.0f} score={score:.0f}")
                print("step,bx,by,pad,target,dx,dy,action,mode,low_steps,reward")
                for row in recent:
                    print(
                        f"{row['step']},{row['ball_x']},{row['ball_y']},"
                        f"{row['paddle_y']},{row['target_y']:.1f},"
                        f"{row['dx']},{row['dy']},"
                        f"{ACTION_NAMES.get(row['action'], row['action'])},"
                        f"{row['mode']},{row['low_serve_steps']},{row['reward']:.0f}"
                    )

            if done:
                print(f"\ndone step={step} score={score:.0f}")
                break
    finally:
        env.close()


if __name__ == "__main__":
    main()
