Apollo3D_SE/world/commons/field.py

from __future__ import annotations

from abc import ABC
from typing import Literal

import numpy as np

from world.commons.field_landmarks import FieldLandmarks

GoalSide = Literal["our", "their", "left", "right"]
Bounds2D = tuple[float, float, float, float]


class Field(ABC):
    FIELD_DIM: tuple[float, float, float]
    LINE_WIDTH: float
    GOAL_DIM: tuple[float, float, float]
    GOALIE_AREA_DIM: tuple[float, float]
    PENALTY_AREA_DIM: tuple[float, float] | None
    PENALTY_SPOT_DISTANCE: float
    CENTER_CIRCLE_RADIUS: float
    DEFAULT_BEAM_POSES: dict[int, tuple[float, float, float]]

    def __init__(self, world):
        from world.world import World  # type hinting

        self.world: World = world
        self.field_landmarks: FieldLandmarks = FieldLandmarks(field=self)

    def _resolve_side(self, side: GoalSide) -> Literal["left", "right"]:
        if side == "left":
            return "left"
        if side == "right":
            return "right"
        is_left = self.world.is_left_team
        if side == "our":
            return "left" if is_left else "right"
        if side == "their":
            return "right" if is_left else "left"
        raise ValueError(f"Unknown field side: {side}")

    def get_width(self) -> float:
        return self.FIELD_DIM[1]

    def get_length(self) -> float:
        return self.FIELD_DIM[0]

    def get_goal_dim(self) -> tuple[float, float, float]:
        return self.GOAL_DIM

    def get_goal_half_width(self) -> float:
        return self.GOAL_DIM[1] / 2.0

    def get_center_circle_radius(self) -> float:
        return self.CENTER_CIRCLE_RADIUS

    def get_goal_position(self, side: GoalSide = "our") -> tuple[float, float]:
        resolved_side = self._resolve_side(side)
        x = -self.get_length() / 2.0 if resolved_side == "left" else self.get_length() / 2.0
        return (x, 0.0)

    def get_our_goal_position(self) -> tuple[float, float]:
        return self.get_goal_position("our")

    def get_their_goal_position(self) -> tuple[float, float]:
        return self.get_goal_position("their")

    def _build_box_bounds(self, depth: float, width: float, side: GoalSide) -> Bounds2D:
        resolved_side = self._resolve_side(side)
        field_half_x = self.get_length() / 2.0
        half_width = width / 2.0

        if resolved_side == "left":
            return (-field_half_x, -field_half_x + depth, -half_width, half_width)
        return (field_half_x - depth, field_half_x, -half_width, half_width)

    def get_goalie_area_bounds(self, side: GoalSide = "our") -> Bounds2D:
        return self._build_box_bounds(
            depth=self.GOALIE_AREA_DIM[0],
            width=self.GOALIE_AREA_DIM[1],
            side=side,
        )

    def get_penalty_area_bounds(self, side: GoalSide = "our") -> Bounds2D:
        if self.PENALTY_AREA_DIM is None:
            raise ValueError(f"{type(self).__name__} does not define a penalty area")

        return self._build_box_bounds(
            depth=self.PENALTY_AREA_DIM[0],
            width=self.PENALTY_AREA_DIM[1],
            side=side,
        )

    def get_penalty_spot(self, side: GoalSide = "our") -> tuple[float, float]:
        resolved_side = self._resolve_side(side)
        x = (self.get_length() / 2.0) - self.PENALTY_SPOT_DISTANCE
        return (-x, 0.0) if resolved_side == "left" else (x, 0.0)

    def _is_inside_bounds(self, pos2d: np.ndarray | tuple[float, float], bounds: Bounds2D) -> bool:
        x, y = float(pos2d[0]), float(pos2d[1])
        min_x, max_x, min_y, max_y = bounds
        return min_x <= x <= max_x and min_y <= y <= max_y

    def is_inside_goalie_area(
        self, pos2d: np.ndarray | tuple[float, float], side: GoalSide = "our"
    ) -> bool:
        return self._is_inside_bounds(pos2d, self.get_goalie_area_bounds(side))

    def is_inside_penalty_area(
        self, pos2d: np.ndarray | tuple[float, float], side: GoalSide = "our"
    ) -> bool:
        return self._is_inside_bounds(pos2d, self.get_penalty_area_bounds(side))

    def is_inside_field(self, pos2d: np.ndarray | tuple[float, float]) -> bool:
        field_half_x = self.get_length() / 2.0
        field_half_y = self.get_width() / 2.0
        return self._is_inside_bounds(pos2d, (-field_half_x, field_half_x, -field_half_y, field_half_y))

    def get_beam_pose(self, number: int, is_left_team: bool = True) -> tuple[float, float, float]:
        try:
            pose = self.DEFAULT_BEAM_POSES[number]
        except KeyError as exc:
            raise KeyError(f"No beam pose configured for player {number} on {type(self).__name__}") from exc

        if is_left_team:
            return pose

        x, y, rotation = pose
        mirrored_rotation = ((rotation + 180.0 + 180.0) % 360.0) - 180.0
        return (-x, -y, mirrored_rotation)

    def get_default_beam_poses(self) -> dict[int, tuple[float, float, float]]:
        return dict(self.DEFAULT_BEAM_POSES)

    def get_canonical_landmarks(self) -> dict[str, np.ndarray]:
        field_half_x = self.get_length() / 2.0
        field_half_y = self.get_width() / 2.0
        goal_half_y = self.get_goal_half_width()
        penalty_marker_x = field_half_x - self.PENALTY_SPOT_DISTANCE
        goalie_area_x = field_half_x - self.GOALIE_AREA_DIM[0]
        goalie_marker_y = self.GOALIE_AREA_DIM[1] / 2.0

        landmarks = {
            "l_luf": np.array([-field_half_x, field_half_y, 0.0]),
            "l_llf": np.array([-field_half_x, -field_half_y, 0.0]),
            "l_ruf": np.array([field_half_x, field_half_y, 0.0]),
            "l_rlf": np.array([field_half_x, -field_half_y, 0.0]),
            "t_cuf": np.array([0.0, field_half_y, 0.0]),
            "t_clf": np.array([0.0, -field_half_y, 0.0]),
            "x_cuc": np.array([0.0, self.CENTER_CIRCLE_RADIUS, 0.0]),
            "x_clc": np.array([0.0, -self.CENTER_CIRCLE_RADIUS, 0.0]),
            "p_lpm": np.array([-penalty_marker_x, 0.0, 0.0]),
            "p_rpm": np.array([penalty_marker_x, 0.0, 0.0]),
            "g_lup": np.array([-field_half_x, goal_half_y, self.GOAL_DIM[2]]),
            "g_llp": np.array([-field_half_x, -goal_half_y, self.GOAL_DIM[2]]),
            "g_rup": np.array([field_half_x, goal_half_y, self.GOAL_DIM[2]]),
            "g_rlp": np.array([field_half_x, -goal_half_y, self.GOAL_DIM[2]]),
            "l_luga": np.array([-goalie_area_x, goalie_marker_y, 0.0]),
            "l_llga": np.array([-goalie_area_x, -goalie_marker_y, 0.0]),
            "l_ruga": np.array([goalie_area_x, goalie_marker_y, 0.0]),
            "l_rlga": np.array([goalie_area_x, -goalie_marker_y, 0.0]),
            "t_luga": np.array([-field_half_x, goalie_marker_y, 0.0]),
            "t_llga": np.array([-field_half_x, -goalie_marker_y, 0.0]),
            "t_ruga": np.array([field_half_x, goalie_marker_y, 0.0]),
            "t_rlga": np.array([field_half_x, -goalie_marker_y, 0.0]),
        }

        if self.PENALTY_AREA_DIM is not None:
            penalty_area_x = field_half_x - self.PENALTY_AREA_DIM[0]
            penalty_marker_y = self.PENALTY_AREA_DIM[1] / 2.0
            landmarks.update(
                {
                    "l_lupa": np.array([-penalty_area_x, penalty_marker_y, 0.0]),
                    "l_llpa": np.array([-penalty_area_x, -penalty_marker_y, 0.0]),
                    "l_rupa": np.array([penalty_area_x, penalty_marker_y, 0.0]),
                    "l_rlpa": np.array([penalty_area_x, -penalty_marker_y, 0.0]),
                    "t_lupa": np.array([-field_half_x, penalty_marker_y, 0.0]),
                    "t_llpa": np.array([-field_half_x, -penalty_marker_y, 0.0]),
                    "t_rupa": np.array([field_half_x, penalty_marker_y, 0.0]),
                    "t_rlpa": np.array([field_half_x, -penalty_marker_y, 0.0]),
                }
            )

        return landmarks


class FIFAField(Field):
    FIELD_DIM = (105.0, 68.0, 40.0)
    LINE_WIDTH = 0.1
    GOAL_DIM = (1.6, 7.32, 2.44)
    GOALIE_AREA_DIM = (5.5, 18.32)
    PENALTY_AREA_DIM = (16.5, 40.32)
    PENALTY_SPOT_DISTANCE = 11.0
    CENTER_CIRCLE_RADIUS = 9.15
    DEFAULT_BEAM_POSES = {
        1: (-50.7, 0.0, 0.0),
        2: (-38.9, 10.9, 0.0),
        3: (-36.8, 4.8, 0.0),
        4: (-36.8, -4.8, 0.0),
        5: (-38.9, -10.9, 0.0),
        6: (-10.2, 0.0, 0.0),
        7: (-18.9, 0.0, 0.0),
    }


class HLAdultField(Field):
    FIELD_DIM = (14.0, 9.0, 40.0)
    LINE_WIDTH = 0.05
    GOAL_DIM = (0.6, 2.6, 1.8)
    GOALIE_AREA_DIM = (1.0, 4.0)
    PENALTY_AREA_DIM = (3.0, 6.0)
    PENALTY_SPOT_DISTANCE = 2.1
    CENTER_CIRCLE_RADIUS = 1.5
    DEFAULT_BEAM_POSES = {
        1: (-5.5, 0.0, 0.0),
        2: (-2.0, -1.5, 0.0),
        3: (-2.0, 1.5, 0.0),
    }


class Soccer7vs7Field(Field):
    FIELD_DIM = (55.0, 36.0, 40.0)
    LINE_WIDTH = 0.1
    GOAL_DIM = (0.84, 3.9, 2.44)
    GOALIE_AREA_DIM = (2.9, 9.6)
    PENALTY_AREA_DIM = (8.6, 21.3)
    PENALTY_SPOT_DISTANCE = 5.8
    CENTER_CIRCLE_RADIUS = 4.79
    DEFAULT_BEAM_POSES = {
        1: (-25.7, 0.0, 0.0),
        2: (-15.7, 5.8, 0.0),
        3: (-13.8, 2.5, 0.0),
        4: (-13.8, -2.5, 0.0),
        5: (-15.7, -5.8, 0.0),
        6: (-5.8, 0.0, 0.0),
        7: (-9.9, 0.0, 0.0),
    }