"""LDS-02 lidar model and the scan-cleaning step shared by the simulator, the dataset pipeline and the real robot.

ROBOTIS LDS-02 (docs.robotis.com, more_info/lds_02): range 0.16 - 8 m, accuracy +-10 mm below 0.3 m and +-3 % from
0.3 to 6 m, ~1 deg resolution, 5 Hz or above. Invalid readings (no return, or closer than 0.16 m) arrive as 0 / inf.

    params = sample_params(rng)                        # once per episode
    raw = lds02_corrupt(true360, rng, params)          # simulator only: true ranges -> what an LDS-02 would report
    cleaner = ScanCleaner(); clean = cleaner(raw, t)   # sim AND robot: invalid -> near-field hold or 3.5, clip

Near-field hold: an invalid ray keeps the smallest valid reading among rays i-2..i+2 in the current scan and the scans
of the last 0.45 s, if that reading is below HOLD_BELOW (an obstacle that came too close to be seen); otherwise it
becomes 3.5 ("no return"), which is how the simulator encoded missing returns during training.
"""
import numpy as np

N_RAYS = 360
MIN_VALID = 0.16          # LDS-02 minimum range
R_MIN, R_MAX = 0.12, 3.5  # training range (JEPA preprocessing and the simulator)
HOLD_BELOW = 0.35
HOLD_WINDOW_S = 0.45
HOLD_NEIGHBOURS = 2


def sample_params(rng):
    """Per-episode sensor parameters (a range of plausible LDS-02 units / surfaces)."""
    return dict(sigma_rel=float(rng.uniform(0.010, 0.020)), sigma_min=0.005, p_drop=float(rng.uniform(0.02, 0.06)),
                p_arc=float(rng.uniform(0.1, 0.4)), offset=int(rng.integers(-2, 3)))


def lds02_corrupt(true_r, rng, p):
    """(360,) true ranges in metres (inf / >= 3.5 = nothing within 3.5 m) -> LDS-02-like readings (0 = invalid)."""
    r = np.asarray(true_r, np.float64).copy()
    far = ~np.isfinite(r) | (r >= R_MAX)
    r = np.where(far, R_MAX, r)
    if p['offset']:
        r, far = np.roll(r, p['offset']), np.roll(far, p['offset'])     # mounting offset, whole degrees
    sigma = np.maximum(p['sigma_min'], p['sigma_rel'] * r)
    r = r + rng.normal(0.0, 1.0, r.shape) * sigma
    r = np.where(far, R_MAX, r)                                          # far stays "far" (clipped at 3.5)
    bad = rng.random(r.shape) < p['p_drop']                              # isolated dropouts
    if rng.random() < p['p_arc']:                                        # 1-3 contiguous dropout arcs (dark / shiny)
        for _ in range(int(rng.integers(1, 4))):
            c, w = int(rng.integers(0, N_RAYS)), int(rng.integers(3, 16))
            bad[(c + np.arange(w)) % N_RAYS] = True
    bad |= r < MIN_VALID                                                 # too close to measure
    return np.where(bad, 0.0, np.minimum(r, R_MAX)).astype(np.float32)


def _neighbour_min(v, k=HOLD_NEIGHBOURS):
    out = v.copy()
    for s in range(1, k + 1):
        out = np.minimum(out, np.minimum(np.roll(v, s, -1), np.roll(v, -s, -1)))
    return out


class ScanCleaner:
    """Stateful per-episode cleaner: call reset() at episode start, then clean = cleaner(raw360, t_seconds).
    hold_window: seconds of memory for the near-field hold (stretch it when the robot is slowed down)."""

    def __init__(self, hold_window=HOLD_WINDOW_S):
        self.hold_window = hold_window
        self.reset()

    def reset(self):
        self.recent = []                                                  # [(t, valid ranges with inf for invalid)]

    def __call__(self, raw, t):
        r = np.asarray(raw, np.float64)
        invalid = ~np.isfinite(r) | (r < R_MIN) | (r <= 0)
        v = np.where(invalid, np.inf, r)
        self.recent = [(tt, vv) for tt, vv in self.recent if t - tt <= self.hold_window] + [(t, v)]
        near = _neighbour_min(np.min(np.stack([vv for _, vv in self.recent]), 0))
        held = np.where(near < HOLD_BELOW, near, R_MAX)
        out = np.where(invalid, held, r)
        return np.clip(out, R_MIN, R_MAX).astype(np.float32), int((invalid & (near < HOLD_BELOW)).sum())