pong_engine.py

"""
Standalone table-tennis physics environment + baseline opponents for Mr. Pong.

This is a faithful port of the reference implementation shipped by the model
author in `inference.py` at https://huggingface.co/fromziro/MrPong (Apache-2.0).
Physics constants, collision handling, observation layout and the baseline
opponent policies are kept 1:1 with the original so that the behaviour shown in
this Space matches the numbers reported on the model card.
"""

import math
import random
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np


@dataclass
class PhysicsConfig:
    table_width: float = 800.0
    table_height: float = 500.0
    paddle_width: float = 14.0
    paddle_height: float = 80.0
    paddle_speed: float = 8.0
    paddle_smoothing: float = 0.70
    ball_radius: float = 8.0
    ball_speed_initial: float = 8.0
    ball_speed_max: float = 16.0
    ball_acceleration: float = 1.035
    frame_skip: int = 3
    max_rally_steps: int = 1500


class StandalonePongEnv:
    """Self-contained table tennis physics environment (port of the author's env)."""

    def __init__(self, phys: Optional[PhysicsConfig] = None, seed: Optional[int] = None):
        self.phys = phys or PhysicsConfig()
        self.rng = random.Random(seed)
        self.ego_paddle_h = self.phys.paddle_height
        self.opp_paddle_h = self.phys.paddle_height
        self.reset()

    def reset(self, serve_direction: Optional[int] = None,
              initial_speed: Optional[float] = None) -> np.ndarray:
        self.ego_y = self.phys.table_height / 2.0
        self.opp_y = self.phys.table_height / 2.0
        self.ego_vy = 0.0
        self.opp_vy = 0.0
        self.prev_ego_action = 0

        self.ball_x = self.phys.table_width / 2.0
        self.ball_y = self.phys.table_height / 2.0

        if serve_direction is None:
            serve_direction = 1 if self.rng.random() < 0.5 else -1

        serve_angle = self.rng.uniform(-math.pi / 7.0, math.pi / 7.0)
        speed = initial_speed or self.phys.ball_speed_initial
        self.ball_vx = serve_direction * speed * math.cos(serve_angle)
        self.ball_vy = speed * math.sin(serve_angle)
        self.rally_count = 0
        self.step_count = 0

        return self.get_ego_observation()

    def _get_action_velocity(self, action: int) -> float:
        if action == 1:
            return -self.phys.paddle_speed
        elif action == 2:
            return self.phys.paddle_speed
        return 0.0

    def physics_substep(self, ego_action: int, opp_action: int) -> Tuple[bool, Dict[str, Any]]:
        info = {"winner": None}
        done = False

        prev_ego_y = self.ego_y
        prev_opp_y = self.opp_y

        ego_target_v = self._get_action_velocity(ego_action)
        opp_target_v = self._get_action_velocity(opp_action)

        alpha = self.phys.paddle_smoothing
        self.ego_vy = alpha * self.ego_vy + (1.0 - alpha) * ego_target_v
        self.opp_vy = alpha * self.opp_vy + (1.0 - alpha) * opp_target_v

        ego_half_h = self.ego_paddle_h / 2.0
        opp_half_h = self.opp_paddle_h / 2.0

        self.ego_y = float(np.clip(self.ego_y + self.ego_vy, ego_half_h,
                                   self.phys.table_height - ego_half_h))
        self.opp_y = float(np.clip(self.opp_y + self.opp_vy, opp_half_h,
                                   self.phys.table_height - opp_half_h))

        prev_ball_x = self.ball_x
        prev_ball_y = self.ball_y
        r = self.phys.ball_radius

        ego_paddle_x = self.phys.paddle_width
        opp_paddle_x = self.phys.table_width - self.phys.paddle_width
        ego_impact_plane = ego_paddle_x + r
        opp_impact_plane = opp_paddle_x - r

        next_ball_x = prev_ball_x + self.ball_vx
        next_ball_y = prev_ball_y + self.ball_vy

        hit_occurred = False

        # Left (ego) paddle — continuous collision detection
        if self.ball_vx < 0 and prev_ball_x >= ego_impact_plane and next_ball_x <= ego_impact_plane:
            t = float(np.clip((prev_ball_x - ego_impact_plane) / max(1e-6, -self.ball_vx), 0.0, 1.0))
            y_ball_at_impact = prev_ball_y + t * self.ball_vy
            y_ego_at_impact = prev_ego_y + t * (self.ego_y - prev_ego_y)

            if abs(y_ball_at_impact - y_ego_at_impact) <= (ego_half_h + r * 0.6):
                hit_occurred = True
                self.rally_count += 1
                offset = float(np.clip((y_ball_at_impact - y_ego_at_impact) / ego_half_h, -1.0, 1.0))
                bounce_angle = offset * (math.pi / 3.0)
                current_speed = math.hypot(self.ball_vx, self.ball_vy)
                new_speed = min(current_speed * self.phys.ball_acceleration, self.phys.ball_speed_max)
                new_vx = new_speed * math.cos(bounce_angle)
                new_vy = new_speed * math.sin(bounce_angle) + 0.25 * self.ego_vy

                rem_dt = 1.0 - t
                self.ball_x = ego_impact_plane + rem_dt * new_vx
                self.ball_y = y_ball_at_impact + rem_dt * new_vy
                self.ball_vx = new_vx
                self.ball_vy = new_vy

        # Right (opponent) paddle
        elif self.ball_vx > 0 and prev_ball_x <= opp_impact_plane and next_ball_x >= opp_impact_plane:
            t = float(np.clip((opp_impact_plane - prev_ball_x) / max(1e-6, self.ball_vx), 0.0, 1.0))
            y_ball_at_impact = prev_ball_y + t * self.ball_vy
            y_opp_at_impact = prev_opp_y + t * (self.opp_y - prev_opp_y)

            if abs(y_ball_at_impact - y_opp_at_impact) <= (opp_half_h + r * 0.6):
                hit_occurred = True
                self.rally_count += 1
                offset = float(np.clip((y_ball_at_impact - y_opp_at_impact) / opp_half_h, -1.0, 1.0))
                bounce_angle = offset * (math.pi / 3.0)
                current_speed = math.hypot(self.ball_vx, self.ball_vy)
                new_speed = min(current_speed * self.phys.ball_acceleration, self.phys.ball_speed_max)
                new_vx = -new_speed * math.cos(bounce_angle)
                new_vy = new_speed * math.sin(bounce_angle) + 0.25 * self.opp_vy

                rem_dt = 1.0 - t
                self.ball_x = opp_impact_plane + rem_dt * new_vx
                self.ball_y = y_ball_at_impact + rem_dt * new_vy
                self.ball_vx = new_vx
                self.ball_vy = new_vy

        if not hit_occurred:
            self.ball_x = next_ball_x
            self.ball_y = next_ball_y

        # Wall collisions
        if self.ball_y - r <= 0:
            self.ball_y = r + abs(r - self.ball_y)
            self.ball_vy = abs(self.ball_vy)
        elif self.ball_y + r >= self.phys.table_height:
            self.ball_y = (self.phys.table_height - r) - abs(self.ball_y + r - self.phys.table_height)
            self.ball_vy = -abs(self.ball_vy)

        # Goal boundaries
        if self.ball_x - r < 0:
            done = True
            info["winner"] = "opponent"
        elif self.ball_x + r > self.phys.table_width:
            done = True
            info["winner"] = "ego"

        self.prev_ego_action = ego_action
        return done, info

    def step(self, ego_action: int, opp_action: int) -> Tuple[np.ndarray, bool, Dict[str, Any]]:
        self.step_count += 1
        done = False
        info = {"winner": None}

        for _ in range(self.phys.frame_skip):
            d, sub_info = self.physics_substep(ego_action, opp_action)
            if d:
                done = True
                info = sub_info
                break

        if not done and self.step_count >= self.phys.max_rally_steps:
            done = True
            info["winner"] = "draw"

        return self.get_ego_observation(), done, info

    def calculate_intercept_y(self, target_x: float, ball_x: float, ball_y: float,
                              ball_vx: float, ball_vy: float) -> float:
        if (target_x > ball_x and ball_vx <= 0) or (target_x < ball_x and ball_vx >= 0):
            return self.phys.table_height / 2.0

        bx, by = float(ball_x), float(ball_y)
        bvx, bvy = float(ball_vx), float(ball_vy)
        h = self.phys.table_height
        r = self.phys.ball_radius

        for _ in range(10):
            dt_x = (target_x - bx) / bvx if bvx != 0 else float("inf")
            if dt_x <= 0:
                break
            if bvy > 0:
                dt_y = (h - r - by) / bvy
            elif bvy < 0:
                dt_y = (r - by) / bvy
            else:
                dt_y = float("inf")

            if dt_x <= dt_y:
                by += bvy * dt_x
                break
            else:
                bx += bvx * dt_y
                by += bvy * dt_y
                bvy = -bvy

        return float(np.clip(by, r, h - r))

    def get_ego_observation(self) -> np.ndarray:
        w, h = self.phys.table_width, self.phys.table_height
        v_max = self.phys.ball_speed_max
        pv_max = self.phys.paddle_speed
        half_h = self.ego_paddle_h / 2.0
        ego_x = self.phys.paddle_width

        pred_intercept_y = self.calculate_intercept_y(ego_x, self.ball_x, self.ball_y,
                                                      self.ball_vx, self.ball_vy)
        rel_pred_y = (pred_intercept_y - self.ego_y) / h
        pred_norm_y = pred_intercept_y / h

        opp_y_norm = self.opp_y / h
        opp_open_top = (self.opp_y - half_h) / h
        opp_open_bottom = (h - (self.opp_y + half_h)) / h
        speed_norm = math.hypot(self.ball_vx, self.ball_vy) / v_max

        return np.array([
            (self.ball_y - self.ego_y) / h,
            (self.ball_x - ego_x) / w,
            self.ball_vx / v_max,
            self.ball_vy / v_max,
            self.ego_y / h,
            self.ego_vy / pv_max,
            (self.opp_y - self.ego_y) / h,
            self.opp_vy / pv_max,
            self.ball_y / h,
            self.ball_x / w,
            rel_pred_y,
            pred_norm_y,
            opp_y_norm,
            opp_open_top,
            opp_open_bottom,
            speed_norm,
        ], dtype=np.float32)

    def get_opp_observation(self) -> np.ndarray:
        w, h = self.phys.table_width, self.phys.table_height
        v_max = self.phys.ball_speed_max
        pv_max = self.phys.paddle_speed
        half_h = self.opp_paddle_h / 2.0
        opp_x = self.phys.table_width - self.phys.paddle_width

        pred_intercept_y = self.calculate_intercept_y(opp_x, self.ball_x, self.ball_y,
                                                      self.ball_vx, self.ball_vy)
        rel_pred_y = (pred_intercept_y - self.opp_y) / h
        pred_norm_y = pred_intercept_y / h

        ego_y_norm = self.ego_y / h
        ego_open_top = (self.ego_y - half_h) / h
        ego_open_bottom = (h - (self.ego_y + half_h)) / h
        speed_norm = math.hypot(self.ball_vx, self.ball_vy) / v_max

        return np.array([
            (self.ball_y - self.opp_y) / h,
            (opp_x - self.ball_x) / w,
            -self.ball_vx / v_max,
            self.ball_vy / v_max,
            self.opp_y / h,
            self.opp_vy / pv_max,
            (self.ego_y - self.opp_y) / h,
            self.ego_vy / pv_max,
            self.ball_y / h,
            (w - self.ball_x) / w,
            rel_pred_y,
            pred_norm_y,
            ego_y_norm,
            ego_open_top,
            ego_open_bottom,
            speed_norm,
        ], dtype=np.float32)


# ==================================================================================
# Baseline opponents (ported from the author's inference.py)
# ==================================================================================

def smooth_aim_action(current_y: float, target_y: float, prev_action: int,
                      deadzone: float = 6.0) -> int:
    diff = target_y - current_y
    if abs(diff) < deadzone:
        return 0
    return 2 if diff > 0 else 1


class RealisticHardOpponent:
    def __init__(self, commit_x_ratio: float = 0.60, rng: Optional[random.Random] = None):
        self.commit_x_ratio = commit_x_ratio
        self.prev_action = 0
        self.perceptual_noise = 0.0
        self.rng = rng or random.Random()

    def act(self, env: StandalonePongEnv) -> int:
        if env.ball_vx <= 0:
            target_y = env.phys.table_height / 2.0
            self.perceptual_noise = self.rng.uniform(-12.0, 12.0)
        elif env.ball_x < env.phys.table_width * self.commit_x_ratio:
            target_y = env.phys.table_height / 2.0 + (env.ball_y - env.phys.table_height / 2.0) * 0.40
        else:
            target_x = env.phys.table_width - env.phys.paddle_width
            exact_y = env.calculate_intercept_y(target_x, env.ball_x, env.ball_y,
                                                env.ball_vx, env.ball_vy)
            target_y = float(np.clip(exact_y + self.perceptual_noise, 8.0,
                                     env.phys.table_height - 8.0))

        action = smooth_aim_action(env.opp_y, target_y, self.prev_action, deadzone=7.0)
        self.prev_action = action
        return action


class MediumOpponent:
    def __init__(self, rng: Optional[random.Random] = None):
        self.prev_action = 0
        self.rng = rng or random.Random()

    def act(self, env: StandalonePongEnv) -> int:
        if env.ball_vx <= 0:
            target_y = env.phys.table_height / 2.0
        else:
            dt = (env.phys.table_width - env.phys.paddle_width - env.ball_x) / max(1.0, env.ball_vx)
            target_y = env.ball_y + env.ball_vy * dt
            target_y = float(np.clip(target_y, 0, env.phys.table_height))

        action = smooth_aim_action(env.opp_y, target_y, self.prev_action, deadzone=14.0)
        self.prev_action = action
        return action


class EasyOpponent:
    def __init__(self, rng: Optional[random.Random] = None):
        self.prev_action = 0
        self.rng = rng or random.Random()

    def act(self, env: StandalonePongEnv) -> int:
        if env.ball_vx <= 0 or env.ball_x < env.phys.table_width * 0.45:
            target_y = env.phys.table_height / 2.0
        else:
            target_y = env.ball_y + self.rng.uniform(-30.0, 30.0)

        action = smooth_aim_action(env.opp_y, target_y, self.prev_action, deadzone=25.0)
        self.prev_action = action
        return action


class ImpossibleHardOpponent:
    def __init__(self, rng: Optional[random.Random] = None):
        self.prev_action = 0
        self.rng = rng or random.Random()

    def act(self, env: StandalonePongEnv) -> int:
        if env.ball_vx <= 0:
            target_y = env.phys.table_height / 2.0
        else:
            target_x = env.phys.table_width - env.phys.paddle_width
            target_y = env.calculate_intercept_y(target_x, env.ball_x, env.ball_y,
                                                 env.ball_vx, env.ball_vy)

        action = smooth_aim_action(env.opp_y, target_y, self.prev_action, deadzone=2.0)
        self.prev_action = action
        return action


class RandomOpponent:
    def __init__(self, rng: Optional[random.Random] = None):
        self.rng = rng or random.Random()

    def act(self, env: StandalonePongEnv) -> int:
        return self.rng.randint(0, 2)


OPPONENTS = {
    "Realistic Hard": RealisticHardOpponent,
    "Medium Logic": MediumOpponent,
    "Easy Logic": EasyOpponent,
    "Impossible Hard": ImpossibleHardOpponent,
    "Random Agent": RandomOpponent,
}