"""Probe Pong-v5 action meanings through RAM byte 51."""

from __future__ import annotations

import argparse
from dataclasses import dataclass

import envpool
import numpy as np


PADDLE_RAM_INDEX = 51


@dataclass
class ActionProbe:
    action: int
    start: int
    end: int
    delta: int
    minimum: int
    maximum: int
    values: list[int]


def step_env(env: object, action: int) -> tuple[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 float(np.asarray(reward).reshape(-1)[0]), bool(np.asarray(done).reshape(-1)[0]), info


def ram_byte(info: dict, index: int) -> int:
    return int(np.asarray(info["ram"]).reshape(-1)[index])


def reset_env(env: object) -> dict:
    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        _obs, info = reset_out
    else:
        _obs, info = reset_out, {}
    return info


def probe_action(action: int, seed: int, warmup_steps: int, hold_steps: int) -> ActionProbe:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        info = reset_env(env)
        for _ in range(warmup_steps):
            _reward, done, info = step_env(env, 0)
            if done:
                info = reset_env(env)

        values = [ram_byte(info, PADDLE_RAM_INDEX)]
        for _ in range(hold_steps):
            _reward, done, info = step_env(env, action)
            values.append(ram_byte(info, PADDLE_RAM_INDEX))
            if done:
                break

        return ActionProbe(
            action=action,
            start=values[0],
            end=values[-1],
            delta=values[-1] - values[0],
            minimum=min(values),
            maximum=max(values),
            values=values,
        )
    finally:
        env.close()


def classify(probes: list[ActionProbe]) -> tuple[list[int], list[int], list[int]]:
    noop = [probe.action for probe in probes if probe.delta == 0]
    up = [probe.action for probe in probes if probe.delta < 0]
    down = [probe.action for probe in probes if probe.delta > 0]
    return noop, up, down


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run each Pong-v5 action from the same seed and compare ram[51]."
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--warmup-steps", type=int, default=30)
    parser.add_argument("--hold-steps", type=int, default=30)
    args = parser.parse_args()

    reference_env = envpool.make("Pong-v5", env_type="gym", num_envs=1)
    try:
        action_count = int(reference_env.action_space.n)
        print(f"envpool: {envpool.__version__}")
        print(f"action space: {reference_env.action_space}")
        print(f"seed: {args.seed}")
        print(f"warmup action: 0")
        print(f"warmup steps: {args.warmup_steps}")
        print(f"hold steps: {args.hold_steps}")
        print(f"tracked RAM byte: ram[{PADDLE_RAM_INDEX}]")
    finally:
        reference_env.close()

    probes = [
        probe_action(
            action=action,
            seed=args.seed,
            warmup_steps=args.warmup_steps,
            hold_steps=args.hold_steps,
        )
        for action in range(action_count)
    ]

    print()
    print("action,start,end,delta,min,max,values")
    for probe in probes:
        print(
            f"{probe.action},"
            f"{probe.start},"
            f"{probe.end},"
            f"{probe.delta},"
            f"{probe.minimum},"
            f"{probe.maximum},"
            f"{probe.values}"
        )

    noop, up, down = classify(probes)
    print()
    print(f"noop candidates: {noop}")
    print(f"up candidates: {up}")
    print(f"down candidates: {down}")


if __name__ == "__main__":
    main()
