ai_mouse/tests/unit/test_postprocess.py

"""Tests for trajectory post-processing primitives."""
from __future__ import annotations

import numpy as np

from ai_mouse._postprocess import gaussian_smooth


def test_gaussian_smooth_preserves_endpoints() -> None:
    x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert result[0] == 1.0
    assert result[-1] == 4.0


def test_gaussian_smooth_short_input_unchanged() -> None:
    x = np.array([1.0, 2.0, 3.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert np.array_equal(result, x)


def test_gaussian_smooth_constant_unchanged() -> None:
    x = np.full(20, 7.5)
    result = gaussian_smooth(x, sigma=1.0)
    assert np.allclose(result, x, atol=1e-6)


from ai_mouse._postprocess import snap_endpoints


def test_snap_endpoints_pins_first_and_last() -> None:
    forward = np.linspace(0.1, 0.9, 16)
    lateral = np.full(16, 0.5)
    f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16)
    assert f[0] == 0.0
    assert l[0] == 0.0
    assert f[-1] == 1.0
    assert l[-1] == 0.0


def test_snap_endpoints_preserves_middle() -> None:
    forward = np.linspace(0.0, 1.0, 16)
    lateral = np.zeros(16)
    f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4)
    # Points before the last n_snap should be unchanged
    assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)


from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start


def test_smooth_start_dampens_lateral() -> None:
    forward = np.linspace(0, 1, 16)
    lateral = np.full(16, 1.0)
    forward[0] = lateral[0] = 0.0  # invariant: snap already done
    _, l = smooth_start(forward.copy(), lateral.copy(), n=4)
    # Lateral at points 1-4 should be < original (dampened)
    assert l[1] < 1.0
    assert l[4] < 1.0
    # Lateral at point 5+ unchanged
    assert l[5] == 1.0


def test_enforce_forward_monotonic_repairs_inversions() -> None:
    f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
    out = enforce_forward_monotonic(f.copy())
    assert np.all(np.diff(out) > 0), out


def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
    f = np.array([-0.1, 0.5, 1.2])
    out = enforce_forward_monotonic(f.copy())
    assert out[0] == 0.0
    assert out[-1] == 1.0


from ai_mouse._postprocess import build_timestamps, resample_arc


def test_resample_arc_identity_when_same_length() -> None:
    pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]])
    out = resample_arc(pts, 4)
    assert np.allclose(out, pts, atol=1e-6)


def test_resample_arc_changes_length() -> None:
    pts = np.array([[float(i), 0.0] for i in range(10)])
    out = resample_arc(pts, 5)
    assert out.shape == (5, 2)
    # Endpoints preserved
    assert np.allclose(out[0], pts[0])
    assert np.allclose(out[-1], pts[-1])


def test_build_timestamps_strictly_increasing() -> None:
    log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0])
    ts = build_timestamps(log_dt, total_duration_ms=200.0)
    assert ts[0] == 0
    assert np.all(np.diff(ts) >= 1)  # at least 1 ms apart


def test_build_timestamps_total_close_to_target() -> None:
    log_dt = np.array([1.0] * 10)
    ts = build_timestamps(log_dt, total_duration_ms=300.0)
    # Last timestamp should be roughly total - one slot
    assert abs(ts[-1] - 270) < 60  # tolerant of clipping


from ai_mouse._postprocess import sample_duration, truncnorm_sample


def test_truncnorm_sample_within_bounds() -> None:
    rng = np.random.default_rng(0)
    samples = [
        truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500)
    ]
    arr = np.array(samples)
    assert arr.min() >= 20.0
    assert arr.max() <= 300.0
    # Mean roughly close to mu
    assert abs(arr.mean() - 80.0) < 5.0


def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None:
    rng = np.random.default_rng(0)
    # mu far outside [low, high] — rejection will fail
    v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng)
    assert 20.0 <= v <= 30.0


def test_sample_duration_uses_correct_bin() -> None:
    dist_dict = {
        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
        "params": [
            {"mu_log": 4.0, "sigma_log": 0.01},  # bin 0: dist < 50
            {"mu_log": 5.0, "sigma_log": 0.01},  # bin 1: 50 <= dist < 100
            {"mu_log": 6.0, "sigma_log": 0.01},  # bin 2: 100 <= dist < 200
        ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5,
    }
    rng = np.random.default_rng(0)
    v = sample_duration(dist_dict, 150.0, rng)
    # exp(6) ~ 403, with tiny sigma we should land near there
    assert 350 < v < 460


from ai_mouse._postprocess import soften_forward


def test_soften_forward_tolerates_small_backtrack() -> None:
    # A 0.01 dip is within the 0.02 tolerance and must survive untouched.
    f = np.array([0.0, 0.30, 0.29, 0.60, 1.0])
    out = soften_forward(f)
    assert np.isclose(out[2], 0.29)


def test_soften_forward_limits_large_backtrack() -> None:
    # A 0.30 dip is noise; it gets pulled up to prev - tol.
    f = np.array([0.0, 0.50, 0.20, 0.70, 1.0])
    out = soften_forward(f)
    assert np.isclose(out[2], 0.50 - 0.02)


def test_soften_forward_allows_moderate_overshoot() -> None:
    # Overshoot past 1.0 is natural; small overshoot survives (compressed
    # but strictly > 1.0).
    f = np.array([0.0, 0.5, 0.9, 1.04, 1.0])
    out = soften_forward(f)
    assert out[3] > 1.0


def test_soften_forward_compresses_extreme_overshoot() -> None:
    # tanh compression: no output value may exceed 1 + overshoot_span.
    f = np.array([0.0, 0.5, 1.30, 1.50, 1.0])
    out = soften_forward(f)
    assert out.max() <= 1.0 + 0.08 + 1e-9
    assert out[2] > 1.0  # still an overshoot, not clipped flat


def test_soften_forward_no_lower_clip() -> None:
    # Small wind-up behind the start is allowed (warp_endpoints pins
    # the first point later; interior may be slightly negative).
    f = np.array([0.0, -0.01, 0.30, 0.70, 1.0])
    out = soften_forward(f)
    assert out[1] < 0.0