"""Probe Pong-v5 RAM bytes for likely ball and paddle state."""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path

import envpool
import numpy as np


WATCHED_BYTES = (13, 14, 49, 50, 51, 54)


@dataclass
class FrameRecord:
    frame: int
    action: int
    reward: float
    done: bool
    ram: list[int]


def planned_action(frame: int) -> int:
    """Use a deterministic sweep to expose player paddle movement."""
    if frame < 80:
        return 0
    if frame < 160:
        return 2
    if frame < 240:
        return 3
    return 0


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 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_list(info: dict) -> list[int]:
    return [int(value) for value in np.asarray(info["ram"]).reshape(-1)[:128]]


def collect_records(seed: int, frames: int) -> list[FrameRecord]:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        reset_env(env)
        records = []
        for frame in range(frames):
            action = planned_action(frame)
            reward, done, info = step_env(env, action)
            records.append(
                FrameRecord(
                    frame=frame,
                    action=action,
                    reward=reward,
                    done=done,
                    ram=ram_list(info),
                )
            )
            if done:
                reset_env(env)
        return records
    finally:
        env.close()


def byte_values(records: list[FrameRecord], index: int) -> list[int]:
    return [record.ram[index] for record in records]


def transition_count(values: list[int]) -> int:
    return sum(1 for before, after in zip(values, values[1:]) if before != after)


def print_watched_ranges(records: list[FrameRecord]) -> None:
    print("watched byte ranges")
    print("byte,min,max,unique,transitions,first,last")
    for index in WATCHED_BYTES:
        values = byte_values(records, index)
        print(
            f"{index},"
            f"{min(values)},"
            f"{max(values)},"
            f"{len(set(values))},"
            f"{transition_count(values)},"
            f"{values[0]},"
            f"{values[-1]}"
        )


def print_frequent_bytes(records: list[FrameRecord], limit: int) -> None:
    rows = []
    for index in range(128):
        values = byte_values(records, index)
        rows.append(
            (
                transition_count(values),
                len(set(values)),
                index,
                min(values),
                max(values),
                values[0],
                values[-1],
            )
        )

    rows.sort(reverse=True)
    print()
    print(f"most frequently changing RAM bytes, top {limit}")
    print("byte,transitions,unique,min,max,first,last")
    for transitions, unique, index, minimum, maximum, first, last in rows[:limit]:
        print(f"{index},{transitions},{unique},{minimum},{maximum},{first},{last}")


def print_frame_samples(records: list[FrameRecord], limit: int) -> None:
    print()
    print(f"first {limit} frame records")
    print("frame,action,reward,done,ram[13],ram[14],ram[49],ram[50],ram[51],ram[54]")
    for record in records[:limit]:
        watched = ",".join(str(record.ram[index]) for index in WATCHED_BYTES)
        print(f"{record.frame},{record.action},{record.reward},{record.done},{watched}")


def print_hypotheses(records: list[FrameRecord]) -> None:
    print()
    print("initial hypotheses")
    print("ram[13]: likely opponent score; it increments on reward -1 events.")
    print("ram[14]: likely player/right score; it increments on reward +1 events.")
    print("ram[49]: likely ball x; it changes steadily across play.")
    print("ram[54]: likely ball y; it changes frequently with vertical ball motion.")
    print("ram[51]: likely player/right paddle y; action 2 lowers it and action 3 raises it.")
    print("ram[50]: likely opponent/left paddle y; it changes independently of our actions.")


def dump_jsonl(records: list[FrameRecord], path: Path) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(
                json.dumps(
                    {
                        "frame": record.frame,
                        "action": record.action,
                        "reward": record.reward,
                        "done": record.done,
                        "ram": record.ram,
                    },
                    separators=(",", ":"),
                )
            )
            handle.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Record 300 Pong-v5 RAM frames and summarize likely state bytes."
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--frames", type=int, default=300)
    parser.add_argument("--top", type=int, default=20)
    parser.add_argument("--sample-frames", type=int, default=12)
    parser.add_argument("--dump-jsonl", type=Path)
    args = parser.parse_args()

    records = collect_records(seed=args.seed, frames=args.frames)
    if args.dump_jsonl is not None:
        dump_jsonl(records, args.dump_jsonl)

    print(f"envpool: {envpool.__version__}")
    print(f"seed: {args.seed}")
    print(f"frames: {len(records)}")
    print("action plan: 0 for frames 0-79, 2 for 80-159, 3 for 160-239, 0 after")
    print()

    print_watched_ranges(records)
    print_frequent_bytes(records, args.top)
    print_frame_samples(records, args.sample_frames)
    if args.dump_jsonl is not None:
        print()
        print(f"wrote frame records: {args.dump_jsonl}")
    print_hypotheses(records)


if __name__ == "__main__":
    main()
