"""RAM state extraction for Pong heuristic policies."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


BALL_X_RAM_INDEX = 49
BALL_Y_RAM_INDEX = 54
SELF_PADDLE_Y_RAM_INDEX = 51
OPPONENT_PADDLE_Y_RAM_INDEX = 50
OPPONENT_SCORE_RAM_INDEX = 13
SELF_SCORE_RAM_INDEX = 14
SCREEN_X_FROM_RAM_OFFSET = -48.5
SCREEN_Y_FROM_RAM_OFFSET = -12.5
PADDLE_SCREEN_Y_FROM_RAM_OFFSET = 2.75


@dataclass(frozen=True)
class PongState:
    ball_x: int
    ball_y: int
    self_paddle_y: int
    opponent_paddle_y: int
    ball_active: bool
    ball_in_right_half: bool
    self_score: int = 0
    opponent_score: int = 0


def ram_from_info(info: dict) -> np.ndarray:
    return np.asarray(info["ram"]).reshape(-1)


def detect_state(info: dict) -> PongState:
    ram = ram_from_info(info)
    ball_x = int(ram[BALL_X_RAM_INDEX])
    ball_y = int(ram[BALL_Y_RAM_INDEX])
    return PongState(
        ball_x=ball_x,
        ball_y=ball_y,
        self_paddle_y=int(ram[SELF_PADDLE_Y_RAM_INDEX]),
        opponent_paddle_y=int(ram[OPPONENT_PADDLE_Y_RAM_INDEX]),
        ball_active=ball_y != 0 and ball_x not in (0, 205),
        ball_in_right_half=ball_x >= 120,
        self_score=int(ram[SELF_SCORE_RAM_INDEX]),
        opponent_score=int(ram[OPPONENT_SCORE_RAM_INDEX]),
    )


def ram_ball_to_screen(ball_x: int, ball_y: int) -> tuple[float, float]:
    return ball_x + SCREEN_X_FROM_RAM_OFFSET, ball_y + SCREEN_Y_FROM_RAM_OFFSET


def ram_y_to_screen(y: int) -> float:
    return y + SCREEN_Y_FROM_RAM_OFFSET


def ram_paddle_y_to_screen(y: int) -> float:
    return y + PADDLE_SCREEN_Y_FROM_RAM_OFFSET
