"""Fast 2D LiDAR simulator for the Gymnasium env.

Raycasts against sheep (discs) and static world geometry (axis-aligned
walls + gate posts) so the env reproduces the false-positive cluster
distribution Webots produces from real 3D geometry.

Returns a range array matching the Webots Lidar device:
180 rays, 140° FOV centred on forward, 12 m max range, 5 mm noise.
See ``protos/ShepherdDog.proto``.
"""

from __future__ import annotations

import math

import numpy as np


# Match protos/ShepherdDog.proto Lidar device.
LIDAR_N_RAYS = 180
LIDAR_FOV = 2.44       # rad ≈ 140°
LIDAR_MAX_RANGE = 12.0
LIDAR_NOISE = 0.005    # m, gaussian std

# Sheep cross-section in the LiDAR plane (horizontal cylinder approx).
SHEEP_RADIUS = 0.30


# --- Static world geometry — mirrors worlds/field.wbt ---

# Vertical walls: (x, y_min, y_max).
_VERTICAL_WALLS = (
    ( 15.0, -15.0,  15.0),  # field east
    (-15.0, -15.0,  15.0),  # field west
    ( 10.0, -22.0, -15.0),  # pen west
    ( 13.0, -22.0, -15.0),  # pen east
)

# Horizontal walls: (y, x_min, x_max). South wall has a 3 m gap at the gate.
_HORIZONTAL_WALLS = (
    ( 15.0, -15.0,  15.0),  # field north
    (-15.0, -15.0,  10.0),  # field south-west of gate
    (-15.0,  13.0,  15.0),  # field south-east of gate
    (-22.0,  10.0,  13.0),  # pen south
)

# Gate posts + field corner pillars, treated as discs at LiDAR height.
_POSTS_XY = np.array([
    ( 10.0, -15.0), ( 13.0, -15.0),
    ( 15.0,  15.0), ( 15.0, -15.0),
    (-15.0,  15.0), (-15.0, -15.0),
], dtype=np.float64)
POST_RADIUS = 0.25


def ray_angles(n: int = LIDAR_N_RAYS, fov: float = LIDAR_FOV) -> np.ndarray:
    """Local-frame ray angles, CCW from forward, sweeping +fov/2 → -fov/2.

    Matches Webots' default Lidar sweep direction.
    """
    return np.linspace(fov / 2.0, -fov / 2.0, n, dtype=np.float64)


_ANGLES = ray_angles()
_COS = np.cos(_ANGLES)
_SIN = np.sin(_ANGLES)


def _raycast_static(
    ox: float, oy: float, cos_w: np.ndarray, sin_w: np.ndarray,
) -> np.ndarray:
    """Per-ray distance to the nearest wall or post hit (∞ if none)."""
    n_rays = cos_w.shape[0]
    best = np.full(n_rays, np.inf, dtype=np.float64)

    EPS = 1e-3
    safe_cos = np.where(np.abs(cos_w) < 1e-9, 1e-9, cos_w)
    safe_sin = np.where(np.abs(sin_w) < 1e-9, 1e-9, sin_w)

    # Vertical walls (x = const)
    for wx, ymin, ymax in _VERTICAL_WALLS:
        t = (wx - ox) / safe_cos
        y_at = oy + t * sin_w
        valid = (t > EPS) & (y_at >= ymin - EPS) & (y_at <= ymax + EPS)
        cand = np.where(valid, t, np.inf)
        np.minimum(best, cand, out=best)

    # Horizontal walls (y = const)
    for wy, xmin, xmax in _HORIZONTAL_WALLS:
        t = (wy - oy) / safe_sin
        x_at = ox + t * cos_w
        valid = (t > EPS) & (x_at >= xmin - EPS) & (x_at <= xmax + EPS)
        cand = np.where(valid, t, np.inf)
        np.minimum(best, cand, out=best)

    # Posts (treat as discs)
    if _POSTS_XY.size:
        px = _POSTS_XY[:, 0] - ox
        py = _POSTS_XY[:, 1] - oy
        t_post = np.outer(px, cos_w) + np.outer(py, sin_w)        # (P, N)
        d2 = (px ** 2 + py ** 2)[:, None]                         # (P, 1)
        perp2 = d2 - t_post ** 2
        R2 = POST_RADIUS ** 2
        hit = (perp2 < R2) & (t_post > 0.0)
        half = np.sqrt(np.clip(R2 - perp2, 0.0, None))
        cand = np.where(hit, t_post - half, np.inf)
        nearest = cand.min(axis=0)
        np.minimum(best, nearest, out=best)

    return best


def simulate_scan(
    dog_x: float, dog_y: float, dog_heading: float,
    sheep_xy: list[tuple[float, float]],
    noise: float = LIDAR_NOISE,
    max_range: float = LIDAR_MAX_RANGE,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """Return a (N,) float32 range array. No-hit entries equal ``max_range``.

    ``sheep_xy`` is every sheep (penned or active) in the scene.
    """
    n_rays = _ANGLES.shape[0]

    ch, sh = math.cos(dog_heading), math.sin(dog_heading)
    cos_w = ch * _COS - sh * _SIN
    sin_w = sh * _COS + ch * _SIN

    # Walls + posts
    best = _raycast_static(dog_x, dog_y, cos_w, sin_w)

    # Sheep discs
    if sheep_xy:
        sx = np.asarray([p[0] for p in sheep_xy], dtype=np.float64) - dog_x
        sy = np.asarray([p[1] for p in sheep_xy], dtype=np.float64) - dog_y
        t = np.outer(sx, cos_w) + np.outer(sy, sin_w)
        s_dist2 = (sx ** 2 + sy ** 2)[:, None]
        perp2 = s_dist2 - t ** 2
        R2 = SHEEP_RADIUS ** 2
        hit = (perp2 < R2) & (t > 0.0)
        half = np.sqrt(np.clip(R2 - perp2, 0.0, None))
        candidate = np.where(hit, t - half, np.inf)
        nearest = candidate.min(axis=0)
        np.minimum(best, nearest, out=best)

    # Entries with no hit stay at inf → clipped to max_range, matching Webots.
    ranges = np.minimum(best, max_range).astype(np.float32)
    return _add_noise(ranges, noise, rng, max_range)


def _add_noise(ranges: np.ndarray, sigma: float,
               rng: np.random.Generator | None, max_range: float) -> np.ndarray:
    if sigma <= 0.0:
        return ranges
    if rng is None:
        rng = np.random.default_rng()
    hit_mask = ranges < max_range - 1e-3
    n_hit = int(hit_mask.sum())
    if n_hit:
        ranges = ranges.copy()
        ranges[hit_mask] += rng.normal(0.0, sigma, size=n_hit).astype(np.float32)
        np.clip(ranges, 0.0, max_range, out=ranges)
    return ranges