#!/usr/bin/env python3
"""Play ALE Pong in a Pygame window with keyboard controls."""

from __future__ import annotations

import argparse
from dataclasses import dataclass

import numpy as np
import pygame
from ale_py import ALEInterface, roms


ACTION_NOOP = 0
ACTION_FIRE = 1
# In ALE's full Atari action set for this Pong ROM, RIGHT moves the right
# paddle up and LEFT moves it down.
ACTION_UP = 3
ACTION_DOWN = 4


@dataclass
class KeyState:
    up: bool = False
    down: bool = False
    fire: bool = False

    def action(self) -> int:
        if self.fire:
            return ACTION_FIRE
        if self.up and not self.down:
            return ACTION_UP
        if self.down and not self.up:
            return ACTION_DOWN
        return ACTION_NOOP

    def action_name(self) -> str:
        if self.fire:
            return "FIRE"
        if self.up and not self.down:
            return "UP"
        if self.down and not self.up:
            return "DOWN"
        return "NOOP"


def set_key_state(keys: KeyState, key: int, value: bool) -> None:
    if key == pygame.K_UP:
        keys.up = value
    elif key == pygame.K_DOWN:
        keys.down = value
    elif key == pygame.K_SPACE:
        keys.fire = value


def current_action(keys: KeyState) -> int:
    if keys.fire:
        return ACTION_FIRE
    if keys.up and not keys.down:
        return ACTION_UP
    if keys.down and not keys.up:
        return ACTION_DOWN
    return ACTION_NOOP


def throttled_action(keys: KeyState, frame_index: int, move_every: int) -> int:
    if keys.fire:
        return ACTION_FIRE
    if frame_index % move_every != 0:
        return ACTION_NOOP
    return current_action(keys)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--scale", type=int, default=3)
    parser.add_argument("--fps", type=int, default=60)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--move-every",
        type=int,
        default=3,
        help="Send one paddle movement every N frames while a key is held.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    ale = ALEInterface()
    ale.setBool("display_screen", False)
    ale.setBool("sound", False)
    ale.setInt("random_seed", args.seed)
    ale.setInt("frame_skip", 1)
    ale.loadROM(str(roms.Pong))

    pygame.init()
    pygame.display.set_caption("ALE Pong")
    screen = pygame.display.set_mode((160 * args.scale, 210 * args.scale))
    clock = pygame.time.Clock()
    font = pygame.font.SysFont("Menlo", 14)
    keys = KeyState()
    score = 0.0
    frame_index = 0
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                elif event.key == pygame.K_r:
                    ale.reset_game()
                    score = 0.0
                else:
                    set_key_state(keys, event.key, True)
            elif event.type == pygame.KEYUP:
                set_key_state(keys, event.key, False)

        if ale.game_over():
            ale.reset_game()
            score = 0.0

        action = throttled_action(keys, frame_index, max(1, args.move_every))
        frame_index += 1
        reward = ale.act(action)
        score += float(reward)

        frame = ale.getScreenRGB()
        frame = np.transpose(frame, (1, 0, 2))
        surface = pygame.surfarray.make_surface(frame)
        surface = pygame.transform.scale(surface, screen.get_size())
        screen.blit(surface, (0, 0))

        label = font.render(
            (
                f"score={score:.0f}  action={keys.action_name()}  "
                f"move_every={args.move_every}  "
                "Up/Down move  Space fire  R reset  Esc quit"
            ),
            True,
            (255, 255, 255),
        )
        screen.blit(label, (8, 8))
        pygame.display.flip()
        clock.tick(args.fps)

    pygame.quit()


if __name__ == "__main__":
    main()
