"""Compare RAM ball bytes with a simple RGB ball detector."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

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 RamPolicy
from pong_state import BALL_X_RAM_INDEX, BALL_Y_RAM_INDEX, SELF_PADDLE_Y_RAM_INDEX


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 connected_components(mask: np.ndarray) -> list[list[tuple[int, int]]]:
    height, width = mask.shape
    seen = np.zeros_like(mask, dtype=bool)
    components = []
    for y in range(height):
        for x in range(width):
            if not mask[y, x] or seen[y, x]:
                continue
            stack = [(x, y)]
            seen[y, x] = True
            pixels = []
            while stack:
                px, py = stack.pop()
                pixels.append((px, py))
                for nx, ny in ((px - 1, py), (px + 1, py), (px, py - 1), (px, py + 1)):
                    if 0 <= nx < width and 0 <= ny < height and mask[ny, nx] and not seen[ny, nx]:
                        seen[ny, nx] = True
                        stack.append((nx, ny))
            components.append(pixels)
    return components


def detect_ball_xy(rgb: np.ndarray) -> tuple[float, float] | None:
    frame = np.asarray(rgb).reshape(-1, 210, 160, 3)[0]
    playfield = frame.copy()
    mask = (
        (playfield[:, :, 0] > 180)
        & (playfield[:, :, 1] > 180)
        & (playfield[:, :, 2] > 180)
    )
    mask[:34, :] = False
    mask[195:, :] = False
    mask[:, :4] = False
    mask[:, 156:] = False

    candidates = []
    for component in connected_components(mask):
        area = len(component)
        if not 2 <= area <= 30:
            continue
        xs = [pixel[0] for pixel in component]
        ys = [pixel[1] for pixel in component]
        width = max(xs) - min(xs) + 1
        height = max(ys) - min(ys) + 1
        if width > 8 or height > 8:
            continue
        candidates.append((area, float(np.mean(xs)), float(np.mean(ys))))

    if not candidates:
        return None
    _area, x, y = max(candidates, key=lambda item: item[0])
    return x, y


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare RGB-detected ball position with RAM bytes.")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--frames", type=int, default=800)
    parser.add_argument("--sample-every", type=int, default=20)
    args = parser.parse_args()

    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=args.seed, render_mode="rgb_array")
    try:
        policy = RamPolicy()
        _obs, info = reset_env(env)
        policy.reset()
        rows = []
        for step in range(1, args.frames + 1):
            action = policy.select_action(info)
            _obs, reward, done, info = step_env(env, action)
            if step % args.sample_every == 0:
                ram = np.asarray(info["ram"]).reshape(-1)
                detected = detect_ball_xy(env.render())
                if detected is not None:
                    visual_x, visual_y = detected
                    rows.append(
                        {
                            "step": step,
                            "ram_x": int(ram[BALL_X_RAM_INDEX]),
                            "ram_y": int(ram[BALL_Y_RAM_INDEX]),
                            "visual_x": visual_x,
                            "visual_y": visual_y,
                            "paddle_y": int(ram[SELF_PADDLE_Y_RAM_INDEX]),
                            "reward": reward,
                        }
                    )
            if done:
                break
    finally:
        env.close()

    print("step,ram_x,ram_y,visual_x,visual_y,paddle_y,reward,dx,dy")
    for row in rows:
        dx = row["visual_x"] - row["ram_x"]
        dy = row["visual_y"] - row["ram_y"]
        print(
            f"{row['step']},{row['ram_x']},{row['ram_y']},"
            f"{row['visual_x']:.1f},{row['visual_y']:.1f},"
            f"{row['paddle_y']},{row['reward']:.0f},{dx:.1f},{dy:.1f}"
        )
    if rows:
        dxs = np.array([row["visual_x"] - row["ram_x"] for row in rows])
        dys = np.array([row["visual_y"] - row["ram_y"] for row in rows])
        print()
        print(f"samples: {len(rows)}")
        print(f"mean visual_x - ram[49]: {float(np.mean(dxs)):.2f}")
        print(f"mean visual_y - ram[54]: {float(np.mean(dys)):.2f}")
        print(f"std visual_x - ram[49]: {float(np.std(dxs)):.2f}")
        print(f"std visual_y - ram[54]: {float(np.std(dys)):.2f}")
    else:
        print("no visual ball candidates detected")


if __name__ == "__main__":
    main()
