From 2231e4e24bcb490cca0c11337a12375ca06b6a14 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:09:22 +0800
Subject: [PATCH] feat(lib): add sample_duration, truncnorm_sample (no scipy)

---
 src/ai_mouse/_postprocess.py   | 41 ++++++++++++++++++++++++++++++++++
 tests/unit/test_postprocess.py | 37 ++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py
index 7ac1c60..3091703 100644
--- a/src/ai_mouse/_postprocess.py
+++ b/src/ai_mouse/_postprocess.py
@@ -137,3 +137,44 @@ def build_timestamps(
         if t_abs[i] <= t_abs[i - 1]:
             t_abs[i] = t_abs[i - 1] + 1.0
     return t_abs
+
+
+def truncnorm_sample(
+    mu: float,
+    sigma: float,
+    low: float,
+    high: float,
+    rng: np.random.Generator,
+    max_tries: int = 32,
+) -> float:
+    """Sample from N(mu, sigma) truncated to [low, high] via rejection.
+
+    Falls back to clipping if rejection fails ``max_tries`` times.
+    """
+    for _ in range(max_tries):
+        v = rng.normal(mu, sigma)
+        if low <= v <= high:
+            return float(v)
+    return float(np.clip(rng.normal(mu, sigma), low, high))
+
+
+def sample_duration(
+    duration_dist: dict,
+    dist: float,
+    rng: np.random.Generator,
+) -> float:
+    """Sample total trajectory duration (ms) for the given pixel distance.
+
+    Uses per-bin log-normal parameters in ``duration_dist``.
+    """
+    bins = duration_dist["bins"]
+    params = duration_dist["params"]
+    bin_idx = len(bins) - 1
+    for i in range(len(bins) - 1):
+        if dist < bins[i + 1]:
+            bin_idx = i
+            break
+    bin_idx = min(bin_idx, len(params) - 1)
+    mu_log = params[bin_idx]["mu_log"]
+    sigma_log = params[bin_idx]["sigma_log"]
+    return float(np.exp(rng.normal(mu_log, sigma_log)))
diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py
index 16e3fbd..42bcd21 100644
--- a/tests/unit/test_postprocess.py
+++ b/tests/unit/test_postprocess.py
@@ -104,3 +104,40 @@ def test_build_timestamps_total_close_to_target() -> None:
     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