game.py

"""
Flappy Bird game engine with headless mode and replay recording.

The game runs at a fixed physics timestep for deterministic behavior.
All gameplay data (states, actions, frames) can be recorded for replay.
"""

import math
import random
import json
import os
from dataclasses import dataclass, asdict
from typing import Optional

# Game constants
SCREEN_WIDTH = 288
SCREEN_HEIGHT = 512
FPS = 60

# Bird constants
BIRD_X = 60
BIRD_WIDTH = 34
BIRD_HEIGHT = 24
GRAVITY = 0.4
FLAP_VELOCITY = -7.0
MAX_FALL_SPEED = 10.0
BIRD_START_Y = 256
BIRD_MIN_X = BIRD_X - 24
BIRD_MAX_X = BIRD_X + 24
WIND_CENTERING = 0.035

# Pipe constants
PIPE_WIDTH = 52
PIPE_GAP = 140  # vertical gap between top and bottom pipes
PIPE_SPEED = 2.5
PIPE_SPAWN_INTERVAL = 100  # frames between pipe spawns
MIN_PIPE_Y = 100  # minimum y for gap center
MAX_PIPE_Y = SCREEN_HEIGHT - 100  # maximum y for gap center

# Ground
GROUND_Y = SCREEN_HEIGHT - 56


@dataclass(frozen=True)
class WorldConfig:
    """Small deterministic game variation used for competition eval."""

    name: str
    world_id: int
    moving_gap_speed: float = 0.0
    wind_strength_x: float = 0.0
    wind_strength_y: float = 0.0
    pipe_speed_jitter: float = 0.0
    pipe_gap_jitter: int = 0


WORLD_CONFIGS = {
    "normal": WorldConfig(name="normal", world_id=0),
    "moving_gap": WorldConfig(name="moving_gap", world_id=1, moving_gap_speed=1.4),
    "wind": WorldConfig(
        name="wind",
        world_id=2,
        wind_strength_x=0.55,
        wind_strength_y=0.12,
    ),
    "variable": WorldConfig(
        name="variable",
        world_id=3,
        moving_gap_speed=1.0,
        pipe_speed_jitter=0.25,
        pipe_gap_jitter=12,
    ),
}


def get_world_config(world: str) -> WorldConfig:
    if world not in WORLD_CONFIGS:
        valid = ", ".join(sorted(WORLD_CONFIGS))
        raise ValueError(f"Unknown world '{world}'. Expected one of: {valid}")
    return WORLD_CONFIGS[world]


@dataclass
class Bird:
    x: float = BIRD_X
    y: float = BIRD_START_Y
    velocity: float = 0.0
    rotation: float = 0.0
    alive: bool = True

    @property
    def rect(self):
        return (self.x - BIRD_WIDTH // 2, self.y - BIRD_HEIGHT // 2,
                BIRD_WIDTH, BIRD_HEIGHT)

    def flap(self, flap_velocity: float = FLAP_VELOCITY):
        self.velocity = flap_velocity

    def update(self, gravity: float = GRAVITY, wind_x: float = 0.0, wind_y: float = 0.0):
        self.velocity = min(self.velocity + gravity + wind_y, MAX_FALL_SPEED)
        self.y += self.velocity
        self.x += wind_x + (BIRD_X - self.x) * WIND_CENTERING
        self.x = max(BIRD_MIN_X, min(BIRD_MAX_X, self.x))
        # Rotation based on velocity
        if self.velocity < 0:
            self.rotation = min(25, -self.velocity * 3)
        else:
            self.rotation = max(-90, -self.velocity * 5)


@dataclass
class Pipe:
    x: float
    gap_center_y: float
    scored: bool = False
    gap_velocity: float = 0.0
    pipe_speed: float = PIPE_SPEED
    pipe_gap: int = PIPE_GAP

    @property
    def top_rect(self):
        top_bottom = self.gap_center_y - self.pipe_gap // 2
        return (self.x, 0, PIPE_WIDTH, top_bottom)

    @property
    def bottom_rect(self):
        bottom_top = self.gap_center_y + self.pipe_gap // 2
        return (self.x, bottom_top, PIPE_WIDTH, SCREEN_HEIGHT - bottom_top)

    def update(self):
        self.x -= self.pipe_speed
        if self.gap_velocity:
            next_y = self.gap_center_y + self.gap_velocity
            if next_y < MIN_PIPE_Y or next_y > MAX_PIPE_Y:
                self.gap_velocity *= -1
                next_y = max(MIN_PIPE_Y, min(MAX_PIPE_Y, next_y))
            self.gap_center_y = next_y


@dataclass
class GameState:
    """Complete game state at a single frame."""
    frame: int
    bird_x: float
    bird_y: float
    bird_velocity: float
    bird_alive: bool
    wind_x: float
    wind_y: float
    pipes: list  # list of (x, gap_center_y, scored)
    score: int
    action: Optional[bool] = None  # True = flap, False = no flap


def rects_collide(r1, r2):
    """Check if two rects (x, y, w, h) overlap."""
    x1, y1, w1, h1 = r1
    x2, y2, w2, h2 = r2
    return (x1 < x2 + w2 and x1 + w1 > x2 and
            y1 < y2 + h2 and y1 + h1 > y2)


class FlappyBirdGame:
    def __init__(self, seed: Optional[int] = None, world: str = "normal"):
        self.seed = seed
        self.world = get_world_config(world)
        self.rng = random.Random(seed)
        self.reset()

    def reset(self):
        """Reset the game to initial state."""
        if self.seed is not None:
            self.rng = random.Random(self.seed)
        self.bird = Bird()
        self.pipes: list[Pipe] = []
        self.score = 0
        self.frame = 0
        self.game_over = False
        self.frames_since_last_pipe = PIPE_SPAWN_INTERVAL - 30  # spawn first pipe sooner
        self.wind_x, self.wind_y = self._calculate_wind()

    def _calculate_wind(self) -> tuple[float, float]:
        if not self.world.wind_strength_x and not self.world.wind_strength_y:
            return 0.0, 0.0
        seed_phase = ((self.seed or 0) % 97) * 0.13
        direction = -1.0 if (self.seed or 0) % 2 else 1.0
        pulse = 0.55 + 0.45 * math.sin(self.frame / 120.0 + seed_phase)
        wind_x = direction * self.world.wind_strength_x * pulse
        wind_y = self.world.wind_strength_y * math.sin(self.frame / 180.0 + seed_phase * 0.7)
        return wind_x, wind_y

    def get_observation(self) -> dict:
        """
        Get the current observation/state for the AI.

        State space design:
        - bird_y: Bird's vertical position (0 = top, SCREEN_HEIGHT = bottom)
        - bird_velocity: Bird's current vertical velocity (negative = going up)
        - next_pipe_dx: Horizontal distance to the next pipe's left edge
        - next_pipe_gap_cy: Y-coordinate of the center of the next pipe gap
        - next_pipe_top_y: Y-coordinate of the bottom edge of the top pipe
        - next_pipe_bottom_y: Y-coordinate of the top edge of the bottom pipe
        - after_pipe_dx: Horizontal distance to the pipe after next
        - after_pipe_gap_cy: Gap center Y of the pipe after next
        - bird_to_gap_dy: Vertical distance from bird to next gap center (positive = below)
        - bird_to_ground: Distance from bird to ground
        - bird_to_ceiling: Distance from bird to ceiling
        """
        # Find the next pipe (first pipe whose right edge is ahead of bird)
        next_pipe = None
        after_pipe = None
        for pipe in self.pipes:
            if pipe.x + PIPE_WIDTH > self.bird.x:
                if next_pipe is None:
                    next_pipe = pipe
                elif after_pipe is None:
                    after_pipe = pipe
                    break

        if next_pipe is None:
            # No pipes on screen yet
            next_pipe_dx = SCREEN_WIDTH - self.bird.x
            next_pipe_gap_cy = SCREEN_HEIGHT / 2
            next_pipe_gap_velocity = 0.0
            next_pipe_speed = PIPE_SPEED
            next_pipe_gap = PIPE_GAP
        else:
            next_pipe_dx = next_pipe.x - self.bird.x
            next_pipe_gap_cy = next_pipe.gap_center_y
            next_pipe_gap_velocity = next_pipe.gap_velocity
            next_pipe_speed = next_pipe.pipe_speed
            next_pipe_gap = next_pipe.pipe_gap

        next_pipe_top_y = next_pipe_gap_cy - next_pipe_gap / 2
        next_pipe_bottom_y = next_pipe_gap_cy + next_pipe_gap / 2

        if after_pipe is None:
            after_pipe_dx = next_pipe_dx + PIPE_SPAWN_INTERVAL * PIPE_SPEED
            after_pipe_gap_cy = SCREEN_HEIGHT / 2
            after_pipe_gap_velocity = 0.0
            after_pipe_speed = PIPE_SPEED
            after_pipe_gap = PIPE_GAP
        else:
            after_pipe_dx = after_pipe.x - self.bird.x
            after_pipe_gap_cy = after_pipe.gap_center_y
            after_pipe_gap_velocity = after_pipe.gap_velocity
            after_pipe_speed = after_pipe.pipe_speed
            after_pipe_gap = after_pipe.pipe_gap

        return {
            "bird_x": self.bird.x,
            "bird_y": self.bird.y,
            "bird_velocity": self.bird.velocity,
            "next_pipe_dx": next_pipe_dx,
            "next_pipe_gap_cy": next_pipe_gap_cy,
            "next_pipe_top_y": next_pipe_top_y,
            "next_pipe_bottom_y": next_pipe_bottom_y,
            "after_pipe_dx": after_pipe_dx,
            "after_pipe_gap_cy": after_pipe_gap_cy,
            "next_pipe_gap_velocity": next_pipe_gap_velocity,
            "after_pipe_gap_velocity": after_pipe_gap_velocity,
            "next_pipe_speed": next_pipe_speed,
            "after_pipe_speed": after_pipe_speed,
            "next_pipe_gap": next_pipe_gap,
            "after_pipe_gap": after_pipe_gap,
            "bird_to_gap_dy": self.bird.y - next_pipe_gap_cy,
            "bird_to_ground": GROUND_Y - self.bird.y,
            "bird_to_ceiling": self.bird.y,
            "pipe_speed": next_pipe_speed,
            "pipe_gap": next_pipe_gap,
            "gravity": GRAVITY,
            "flap_velocity": FLAP_VELOCITY,
            "wind_x": self.wind_x,
            "wind_y": self.wind_y,
            "world_type": self.world.name,
            "world_id": self.world.world_id,
        }

    def step(self, flap: bool = False) -> tuple[dict, float, bool]:
        """
        Advance the game by one frame.

        Args:
            flap: Whether the bird should flap this frame.

        Returns:
            (observation, reward, done)
        """
        if self.game_over:
            return self.get_observation(), 0.0, True

        self.frame += 1
        reward = 0.1  # small reward for surviving each frame

        # Bird action
        if flap:
            self.bird.flap(FLAP_VELOCITY)

        # Update bird
        self.wind_x, self.wind_y = self._calculate_wind()
        self.bird.update(GRAVITY, self.wind_x, self.wind_y)

        # Spawn pipes
        self.frames_since_last_pipe += 1
        if self.frames_since_last_pipe >= PIPE_SPAWN_INTERVAL:
            gap_y = self.rng.randint(MIN_PIPE_Y, MAX_PIPE_Y)
            gap_velocity = 0.0
            if self.world.moving_gap_speed:
                gap_velocity = self.world.moving_gap_speed * self.rng.choice([-1, 1])
            pipe_speed = PIPE_SPEED
            if self.world.pipe_speed_jitter:
                pipe_speed += self.rng.uniform(
                    -self.world.pipe_speed_jitter,
                    self.world.pipe_speed_jitter,
                )
            pipe_gap = PIPE_GAP
            if self.world.pipe_gap_jitter:
                pipe_gap += self.rng.randint(
                    -self.world.pipe_gap_jitter,
                    self.world.pipe_gap_jitter,
                )
            self.pipes.append(Pipe(
                x=SCREEN_WIDTH,
                gap_center_y=gap_y,
                gap_velocity=gap_velocity,
                pipe_speed=pipe_speed,
                pipe_gap=pipe_gap,
            ))
            self.frames_since_last_pipe = 0

        # Update pipes
        for pipe in self.pipes:
            pipe.update()

        # Remove off-screen pipes
        self.pipes = [p for p in self.pipes if p.x + PIPE_WIDTH > -10]

        # Scoring
        for pipe in self.pipes:
            if not pipe.scored and pipe.x + PIPE_WIDTH < self.bird.x:
                pipe.scored = True
                self.score += 1
                reward = 1.0  # bonus for passing a pipe

        # Collision detection
        bird_rect = self.bird.rect
        # Ground/ceiling
        if self.bird.y + BIRD_HEIGHT // 2 >= GROUND_Y or self.bird.y - BIRD_HEIGHT // 2 <= 0:
            self.game_over = True
            self.bird.alive = False
            reward = -5.0

        # Pipe collision
        if not self.game_over:
            for pipe in self.pipes:
                if (rects_collide(bird_rect, pipe.top_rect) or
                        rects_collide(bird_rect, pipe.bottom_rect)):
                    self.game_over = True
                    self.bird.alive = False
                    reward = -5.0
                    break

        return self.get_observation(), reward, self.game_over

    def get_state(self, action: Optional[bool] = None) -> GameState:
        """Capture the full game state for replay."""
        return GameState(
            frame=self.frame,
            bird_x=self.bird.x,
            bird_y=self.bird.y,
            bird_velocity=self.bird.velocity,
            bird_alive=self.bird.alive,
            wind_x=self.wind_x,
            wind_y=self.wind_y,
            pipes=[
                (p.x, p.gap_center_y, p.scored, p.gap_velocity, p.pipe_gap, p.pipe_speed)
                for p in self.pipes
            ],
            score=self.score,
            action=action,
        )


class GameRecorder:
    """Records game states for replay and video generation."""

    def __init__(self, output_dir: str = "replays"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.states: list[dict] = []
        self.metadata: dict = {}

    def record_state(self, state: GameState):
        self.states.append(asdict(state))

    def save(self, filename: str, metadata: Optional[dict] = None):
        """Save the replay to a JSON file."""
        data = {
            "metadata": metadata or {},
            "game_constants": {
                "screen_width": SCREEN_WIDTH,
                "screen_height": SCREEN_HEIGHT,
                "bird_x": BIRD_X,
                "bird_min_x": BIRD_MIN_X,
                "bird_max_x": BIRD_MAX_X,
                "bird_width": BIRD_WIDTH,
                "bird_height": BIRD_HEIGHT,
                "pipe_width": PIPE_WIDTH,
                "pipe_gap": PIPE_GAP,
                "pipe_speed": PIPE_SPEED,
                "ground_y": GROUND_Y,
            },
            "total_frames": len(self.states),
            "final_score": self.states[-1]["score"] if self.states else 0,
            "states": self.states,
        }
        filepath = os.path.join(self.output_dir, filename)
        with open(filepath, "w") as f:
            json.dump(data, f)
        return filepath

    def reset(self):
        self.states = []