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

Raycasts against:
  * **Sheep** — discs of radius ``SHEEP_RADIUS``.
  * **Static world geometry** — axis-aligned wall segments and gate
    posts taken from ``worlds/field.wbt``. Without these, demos
    collected in-env would never include the false-positive clusters
    Webots produces from the stone walls and gate-post boxes, and the
    BC student trained on those demos collapses on deployment.

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

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 modelled as a vertical cylinder; this is the horizontal-section
# radius the LiDAR plane intersects. Tuned to the proto sheep (~0.45 m
# body length). The exact value is not load-bearing — the perception
# clusterer is range-tolerant.
SHEEP_RADIUS = 0.30


# ---------------------------------------------------------------------------
# Static world geometry — must match worlds/field.wbt
# ---------------------------------------------------------------------------

# Vertical walls: (x, y_min, y_max). Field east/west walls and the two
# pen side walls are visible through the open gate.
_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 is split by the 3 m
# gate at x ∈ [10, 13]; the pen south wall closes the back of the pen.
_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 and field corner pillars treated as vertical cylinders at
# LiDAR height. Radius 0.25 m comes from the 0.44 × 0.44 m boxes in the
# wbt — close enough to a circular cross-section for this purpose.
_POSTS_XY = np.array([
    ( 10.0, -15.0),  # west gate post
    ( 13.0, -15.0),  # east gate post
    ( 15.0,  15.0),  # NE field corner
    ( 15.0, -15.0),  # SE field corner
    (-15.0,  15.0),  # NW field corner
    (-15.0, -15.0),  # SW field corner
], 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, sweeping from +fov/2 to -fov/2.

    Convention: angle is measured CCW from the dog's forward axis. Ray 0
    points to the dog's left, last ray to the right. Webots' default
    Lidar sweep matches this.
    """
    return np.linspace(fov / 2.0, -fov / 2.0, n, dtype=np.float64)


# Cached so we don't rebuild every step.
_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 nearest wall or post hit (∞ if none).

    Walls are axis-aligned line segments; for each ray we compute t at
    which it crosses the wall's constant-coord plane and check the
    other coord lies in the segment. Posts are circles; same disc
    intersection as for sheep.
    """
    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 the list of (x, y) world positions of every sheep in
    the scene (penned and active). Static world geometry (walls and
    posts) is also raycast so demos contain the same false-positive
    clusters Webots produces.
    """
    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)

    # Clip to LIDAR_MAX_RANGE; entries that never got a hit stay at inf
    # → clipped down to max_range like the real Webots device.
    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