"""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 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


from ai_mouse._postprocess import damp_start


def test_damp_start_dampens_early_lateral() -> None:
    lat = np.full(16, 1.0)
    out = damp_start(lat, n=4)
    assert out[1] < out[2] < out[3] < out[4] < 1.0  # monotone ramp
    assert np.all(out[5:] == 1.0)  # untouched past n


def test_damp_start_no_release_jump() -> None:
    # The weight at i=n must be close to 1 (continuous release):
    # smoothstep(4/5) = 0.896, vs the old linear 4/5 = 0.8.
    lat = np.full(16, 1.0)
    out = damp_start(lat, n=4)
    assert out[4] > 0.85


def test_damp_start_short_input_safe() -> None:
    lat = np.array([0.0, 0.5, 0.3])
    out = damp_start(lat, n=4)  # n capped to len//4 = 0 → no-op
    assert np.array_equal(out, lat)


from ai_mouse._postprocess import warp_endpoints


def test_warp_endpoints_exact_pin() -> None:
    f = np.linspace(0.05, 1.10, 32)
    l = np.linspace(0.03, -0.07, 32)
    fo, lo = warp_endpoints(f, l)
    assert fo[0] == 0.0 and lo[0] == 0.0
    assert fo[-1] == 1.0 and lo[-1] == 0.0


def test_warp_endpoints_identity_when_already_pinned() -> None:
    f = np.linspace(0.0, 1.0, 32)
    l = np.sin(np.linspace(0, np.pi, 32)) * 0.1
    l[0] = l[-1] = 0.0
    fo, lo = warp_endpoints(f.copy(), l.copy())
    assert np.allclose(fo, f, atol=1e-12)
    assert np.allclose(lo, l, atol=1e-12)


def test_warp_endpoints_preserves_smoothness() -> None:
    # Correcting a smooth curve must not introduce sharp local bends:
    # the warp adds a smoothstep-weighted offset, so the second
    # difference (discrete curvature proxy) stays small.
    f = np.linspace(0.02, 1.08, 32)
    l = np.full(32, 0.05)
    fo, lo = warp_endpoints(f, l)
    assert np.abs(np.diff(lo, 2)).max() < 0.01
    assert np.abs(np.diff(fo, 2)).max() < 0.01


def test_warp_endpoints_correction_local_to_each_end() -> None:
    # A start-only residual should barely move the last quarter.
    # (smoothstep weight at i=24/31 is ~0.13, so 0.08 residual leaves
    # ~0.010 there — threshold 0.02 gives margin without losing meaning)
    f = np.linspace(0.0, 1.0, 32) + 0.0
    l = np.zeros(32)
    f[0] = 0.08  # start residual only
    fo, _ = warp_endpoints(f, l)
    assert np.abs(fo[24:] - f[24:]).max() < 0.02