From b93d240641b2a710f0f1775043aa97706ff3adaa Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:08:39 +0800
Subject: [PATCH] feat(lib): add resample_arc, build_timestamps

---
 src/ai_mouse/_postprocess.py   | 47 ++++++++++++++++++++++++++++++++++
 tests/unit/test_postprocess.py | 32 +++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py
index ac2b6e3..7ac1c60 100644
--- a/src/ai_mouse/_postprocess.py
+++ b/src/ai_mouse/_postprocess.py
@@ -90,3 +90,50 @@ def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
     forward[0] = 0.0
     forward[-1] = 1.0
     return forward
+
+
+def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
+    """Resample a 2-D polyline to ``n_points`` along cumulative arc length."""
+    arc = np.concatenate(
+        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
+    )
+    s_new = np.linspace(0, arc[-1], n_points)
+    return np.stack(
+        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
+        axis=1,
+    )
+
+
+def build_timestamps(
+    log_dt: np.ndarray,
+    total_duration_ms: float,
+    dt_clip: tuple[float, float] = (2.0, 150.0),
+) -> np.ndarray:
+    """Convert per-step log_dt + total duration to cumulative ms timestamps.
+
+    Args:
+        log_dt: (N,) array of natural-log step intervals.
+        total_duration_ms: target total span. The output is scaled so the
+            sum approximately matches this (modulo dt_clip).
+        dt_clip: (min, max) per-step clamp in milliseconds.
+
+    Returns:
+        (N,) integer-rounded cumulative timestamps starting at 0,
+        strictly increasing.
+    """
+    n = len(log_dt)
+    dt_raw = np.clip(np.exp(log_dt), 0.0, None)
+    dt_sum = dt_raw.sum()
+    if dt_sum > 1e-6:
+        scale = total_duration_ms / dt_sum
+    else:
+        scale = total_duration_ms / max(n, 1)
+    dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1])
+
+    t_abs = np.cumsum(dt_ms)
+    t_abs = np.concatenate([[0.0], t_abs[:-1]])
+
+    for i in range(1, n):
+        if t_abs[i] <= t_abs[i - 1]:
+            t_abs[i] = t_abs[i - 1] + 1.0
+    return t_abs
diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py
index 0ec1439..16e3fbd 100644
--- a/tests/unit/test_postprocess.py
+++ b/tests/unit/test_postprocess.py
@@ -72,3 +72,35 @@ def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
     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