diff --git a/ai_mouse/data_adapters/balabit.py b/ai_mouse/data_adapters/balabit.py index 5f422f3..da10070 100644 --- a/ai_mouse/data_adapters/balabit.py +++ b/ai_mouse/data_adapters/balabit.py @@ -124,9 +124,72 @@ def filter_segments( min_dist: int, max_span_ms: int, max_gap_ms: int, + coord_max: int = 5000, ) -> list[Segment]: - """Stub — implemented in Task 5.""" - raise NotImplementedError + """Drop segments that fail any quality check. + + A segment is dropped if any of these are true: + - len(events) < min_events + - Euclidean dist(events[0], (click_x, click_y)) < min_dist + - events[-1].t_ms - events[0].t_ms > max_span_ms + - any adjacent Move pair has dt > max_gap_ms (sampling drop-out) + - any coord (start/end/click) outside [0, coord_max] + - total arc length < min_dist (high-frequency jitter only) + + Args: + segments: Candidate segments. + min_events: Minimum number of Move events. + min_dist: Minimum start→click pixel distance AND minimum total arc length. + max_span_ms: Maximum time span of the segment (events[-1] - events[0]). + max_gap_ms: Maximum allowed gap between adjacent Move events. + coord_max: Maximum allowed pixel coordinate value (5000 catches multi-monitor anomalies). + + Returns: + Filtered list, original order preserved. + """ + import math + + keep: list[Segment] = [] + for seg in segments: + if len(seg.events) < min_events: + continue + + sx, sy = seg.events[0].x, seg.events[0].y + ex, ey = seg.click_x, seg.click_y + + # Coord range check + if any(c < 0 or c > coord_max for c in (sx, sy, ex, ey)): + continue + + # Endpoint distance + dist = math.hypot(ex - sx, ey - sy) + if dist < min_dist: + continue + + # Time span + span = seg.events[-1].t_ms - seg.events[0].t_ms + if span > max_span_ms: + continue + + # Gap check + total arc length + total_arc = 0.0 + bad_gap = False + for i in range(1, len(seg.events)): + dt = seg.events[i].t_ms - seg.events[i - 1].t_ms + if dt > max_gap_ms: + bad_gap = True + break + dx = seg.events[i].x - seg.events[i - 1].x + dy = seg.events[i].y - seg.events[i - 1].y + total_arc += math.hypot(dx, dy) + if bad_gap: + continue + if total_arc < min_dist: + continue + + keep.append(seg) + + return keep def process_session( diff --git a/tests/test_balabit_adapter.py b/tests/test_balabit_adapter.py index 1bbe8ac..d6bf1f0 100644 --- a/tests/test_balabit_adapter.py +++ b/tests/test_balabit_adapter.py @@ -185,3 +185,79 @@ class TestSegmentByClicks: # Drag event should not appear in seg.events assert all(e.state == "Move" for e in segments[0].events) assert len(segments[0].events) == 2 + + +class TestFilterSegments: + def _seg(self, events_data: list[tuple[int, int, int]], click=(500, 500), session="s"): + """events_data: list of (t_ms, x, y) tuples.""" + from ai_mouse.data_adapters.balabit import MouseEvent, Segment + events = [MouseEvent(t_ms=t, button="NoButton", state="Move", x=x, y=y) + for (t, x, y) in events_data] + click_t = events[-1].t_ms + 50 if events else 100 + return Segment(events=events, click_x=click[0], click_y=click[1], + click_t_ms=click_t, session_id=session) + + def test_drops_segment_with_too_few_events(self): + from ai_mouse.data_adapters.balabit import filter_segments + seg = self._seg([(0, 100, 100), (50, 105, 105)]) # 2 events, min=5 + result = filter_segments([seg], min_events=5, min_dist=10, + max_span_ms=5000, max_gap_ms=200) + assert result == [] + + def test_drops_segment_with_short_distance(self): + from ai_mouse.data_adapters.balabit import filter_segments + # Start (100,100), end click=(105,100) → dist=5, min_dist=50 + events = [(i*10, 100+i, 100) for i in range(10)] + seg = self._seg(events, click=(105, 100)) + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=200) + assert result == [] + + def test_drops_segment_with_too_long_span(self): + from ai_mouse.data_adapters.balabit import filter_segments + # Span 6000ms, max_span=5000 + events = [(0, 100, 100), (2000, 200, 200), (4000, 300, 300), (6000, 400, 400), + (6010, 410, 400), (6020, 420, 400)] + seg = self._seg(events, click=(500, 500)) + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=10000) + assert result == [] + + def test_drops_segment_with_gap(self): + from ai_mouse.data_adapters.balabit import filter_segments + # Gap of 500ms between events 2 and 3, max_gap=200 + events = [(0, 100, 100), (50, 110, 110), (100, 120, 120), + (600, 200, 200), (650, 210, 210), (700, 220, 220)] + seg = self._seg(events, click=(300, 300)) + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=200) + assert result == [] + + def test_drops_segment_with_out_of_range_coords(self): + from ai_mouse.data_adapters.balabit import filter_segments + # x=10000 out of range + events = [(i*10, 10000, 100) for i in range(6)] + seg = self._seg(events, click=(10100, 100)) + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=200) + assert result == [] + + def test_drops_segment_with_short_arc_length(self): + """Total arc < 50 even though endpoints are far → high-frequency jitter only.""" + from ai_mouse.data_adapters.balabit import filter_segments + # Tiny back-and-forth with click 100px away — unrealistic, drop it + events = [(i*10, 100 + (i % 2), 100) for i in range(10)] + seg = self._seg(events, click=(200, 100)) + # arc length = sum of |Δp| ≈ 9 (alternating ±1) which is < 50 + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=200) + assert result == [] + + def test_keeps_valid_segment(self): + from ai_mouse.data_adapters.balabit import filter_segments + # Smooth 100→500 px straight line, 10 events, span 500ms + events = [(i*50, 100 + i*40, 100) for i in range(10)] + seg = self._seg(events, click=(500, 100)) + result = filter_segments([seg], min_events=5, min_dist=50, + max_span_ms=5000, max_gap_ms=200) + assert len(result) == 1