"""Inspect the minimal EnvPool Pong surface for heuristic learning."""

from __future__ import annotations

import argparse

import envpool
import numpy as np


def flatten_prefix(values: object, count: int) -> list[int]:
    """Return the first count values from an array-like object as ints."""
    return [int(x) for x in np.asarray(values).reshape(-1)[:count]]


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Create Pong-v5, reset once, and print a short RAM trace."
    )
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--action", type=int, default=0)
    args = parser.parse_args()

    env = envpool.make("Pong-v5", env_type="gym", num_envs=1)

    print(f"envpool: {envpool.__version__}")
    print(f"env type: {type(env).__name__}")
    print(f"action space: {env.action_space}")
    print(f"observation space: {env.observation_space}")

    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        obs, info = reset_out
    else:
        obs, info = reset_out, {}

    print(f"reset obs shape: {getattr(obs, 'shape', None)}")
    print(f"reset obs dtype: {getattr(obs, 'dtype', None)}")
    print(f"reset info keys: {sorted(info.keys())}")

    if "ram" in info:
        ram = np.asarray(info["ram"])
        print(f"reset ram shape: {ram.shape}")
        print(f"reset ram dtype: {ram.dtype}")
        print(f"reset ram[0:64]: {flatten_prefix(ram, 64)}")

    action = np.array([args.action], dtype=np.int32)
    for step in range(1, args.steps + 1):
        step_out = env.step(action)
        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

        fields = [
            f"step {step:02d}",
            f"action={args.action}",
            f"reward={np.asarray(reward).tolist()}",
            f"done={np.asarray(done).tolist()}",
        ]
        if "ram" in info:
            fields.append(f"ram[0:16]={flatten_prefix(info['ram'], 16)}")
        print(" ".join(fields))

    env.close()


if __name__ == "__main__":
    main()
