# sites/ai_mouse/ai_mouse/_collector.py from __future__ import annotations import json import math import random from enum import Enum, auto from pathlib import Path class CollectorState(Enum): IDLE = auto() HOVER_A = auto() RECORDING = auto() class Collector: """Manages A→B mouse movement trace collection state and persistence. The state-machine methods (_on_mouse_motion, _on_mouseup, _on_skip) are public so they can be called by the web API without a display. A/B positions are exposed as .a_pos / .b_pos attributes. """ POINT_RADIUS = 15 # pixels — must be inside this to hover/click DWELL_MS = 200 # milliseconds to dwell inside A before recording starts def __init__( self, count: int, dist_min: int, dist_max: int, output_path: Path, screen_size: tuple[int, int] = (800, 600), ): self.count = count self.dist_min = dist_min self.dist_max = dist_max self.output_path = Path(output_path) self.screen_w, self.screen_h = screen_size self.collected = 0 self.state = CollectorState.IDLE self._buffer: list[dict] = [] self._hover_enter_t: int = 0 self._record_start_t: int = 0 self.a_pos, self.b_pos = self._new_ab() # ------------------------------------------------------------------ # State machine (called by web API) # ------------------------------------------------------------------ def _on_mouse_motion(self, mx: int, my: int, t: int) -> None: """Handle a MOUSEMOTION event at pixel (mx, my), time t ms from start.""" if self.state == CollectorState.IDLE: if self._inside(mx, my, self.a_pos): self.state = CollectorState.HOVER_A self._hover_enter_t = t elif self.state == CollectorState.HOVER_A: if not self._inside(mx, my, self.a_pos): self.state = CollectorState.IDLE elif t - self._hover_enter_t >= self.DWELL_MS: self.state = CollectorState.RECORDING self._record_start_t = t self._buffer = [{"type": "move", "x": mx, "y": my, "t": 0}] elif self.state == CollectorState.RECORDING: rel_t = t - self._record_start_t self._buffer.append({"type": "move", "x": mx, "y": my, "t": rel_t}) def _on_mousedown(self, mx: int, my: int, t: int) -> None: """Handle a MOUSEBUTTONDOWN event.""" if self.state == CollectorState.RECORDING: rel_t = t - self._record_start_t self._buffer.append({"type": "down", "x": mx, "y": my, "t": rel_t}) def _on_mouseup(self, mx: int, my: int, t: int) -> None: """Handle a MOUSEBUTTONUP event.""" if self.state == CollectorState.RECORDING: rel_t = t - self._record_start_t self._buffer.append({"type": "up", "x": mx, "y": my, "t": rel_t}) if self._inside(mx, my, self.b_pos): self._save_trace() self.collected += 1 if self.collected < self.count: self.a_pos, self.b_pos = self._new_ab() self.state = CollectorState.IDLE else: # Click outside B — discard buffer and regenerate self._on_skip() def _on_skip(self) -> None: """Handle ESC/skip — discard current buffer, regenerate A/B.""" self._buffer = [] self.state = CollectorState.IDLE self.a_pos, self.b_pos = self._new_ab() # ------------------------------------------------------------------ # Persistence # ------------------------------------------------------------------ def _save_trace(self) -> None: dist = self._dist(self.a_pos, self.b_pos) angle = math.degrees( math.atan2( self.b_pos[1] - self.a_pos[1], self.b_pos[0] - self.a_pos[0], ) ) trace = { "meta": { "start": list(self.a_pos), "end": list(self.b_pos), "dist": round(dist), "angle": round(angle, 1), }, "events": list(self._buffer), } self.output_path.parent.mkdir(parents=True, exist_ok=True) with self.output_path.open("a", encoding="utf-8") as f: f.write(json.dumps(trace, ensure_ascii=False) + "\n") self._buffer = [] # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ def _inside(self, mx: int, my: int, pos: tuple[int, int]) -> bool: return self._dist((mx, my), pos) <= self.POINT_RADIUS @staticmethod def _dist(a: tuple[int, int], b: tuple[int, int]) -> float: return math.hypot(a[0] - b[0], a[1] - b[1]) def _new_ab(self) -> tuple[tuple[int, int], tuple[int, int]]: """Generate a new random A→B pair within distance constraints. Strategy: clamp dist_min/dist_max to the canvas diagonal. When the required distance is large relative to the canvas, bias A towards edges and corners so that long-distance B positions become reachable. The fallback randomly picks from the four corner pairs with jitter to ensure variety even in the degenerate case. """ margin = self.POINT_RADIUS + 5 x_lo, x_hi = margin, self.screen_w - margin y_lo, y_hi = margin, self.screen_h - margin w_inner, h_inner = x_hi - x_lo, y_hi - y_lo max_possible = int(math.hypot(w_inner, h_inner)) eff_max = min(self.dist_max, max_possible) eff_min = min(self.dist_min, eff_max) # Determine how "tight" the distance requirement is relative to canvas. # When ratio > 0.7, purely random A rarely works — bias towards edges. tightness = eff_min / max_possible if max_possible > 0 else 1.0 for _ in range(500): if tightness > 0.7: # Bias A towards edges/corners: pick from a ring near the border side = random.choice(["top", "bottom", "left", "right"]) edge_band = max(int(w_inner * 0.15), 1) if side == "top": ax = random.randint(x_lo, x_hi) ay = random.randint(y_lo, y_lo + edge_band) elif side == "bottom": ax = random.randint(x_lo, x_hi) ay = random.randint(y_hi - edge_band, y_hi) elif side == "left": ax = random.randint(x_lo, x_lo + edge_band) ay = random.randint(y_lo, y_hi) else: ax = random.randint(x_hi - edge_band, x_hi) ay = random.randint(y_lo, y_hi) else: ax = random.randint(x_lo, x_hi) ay = random.randint(y_lo, y_hi) # Compute the farthest reachable distance from (ax, ay) within bounds reach = max( math.hypot(ax - x_lo, ay - y_lo), math.hypot(ax - x_hi, ay - y_lo), math.hypot(ax - x_lo, ay - y_hi), math.hypot(ax - x_hi, ay - y_hi), ) if reach < eff_min: continue local_max = min(eff_max, int(reach)) # Try several angles from this A for _ in range(30): angle = random.uniform(0, 2 * math.pi) dist = random.randint(eff_min, local_max) bx = int(ax + dist * math.cos(angle)) by = int(ay + dist * math.sin(angle)) if x_lo <= bx <= x_hi and y_lo <= by <= y_hi: return (ax, ay), (bx, by) # Fallback: pick a random corner pair with jitter for variety corners = [(x_lo, y_lo), (x_hi, y_lo), (x_lo, y_hi), (x_hi, y_hi)] pairs = [(corners[i], corners[j]) for i in range(4) for j in range(i + 1, 4) if self._dist(corners[i], corners[j]) >= eff_min] if not pairs: # All pairs too short — pick the longest pair pairs = [(corners[i], corners[j]) for i in range(4) for j in range(i + 1, 4)] pairs.sort(key=lambda p: self._dist(p[0], p[1]), reverse=True) pairs = pairs[:1] ca, cb = random.choice(pairs) # Add jitter so it's not identical each time jitter = max(margin, int(min(w_inner, h_inner) * 0.08)) ax = ca[0] + random.randint(-jitter, jitter) ay = ca[1] + random.randint(-jitter, jitter) bx = cb[0] + random.randint(-jitter, jitter) by = cb[1] + random.randint(-jitter, jitter) # Clamp back into bounds ax = max(x_lo, min(x_hi, ax)) ay = max(y_lo, min(y_hi, ay)) bx = max(x_lo, min(x_hi, bx)) by = max(y_lo, min(y_hi, by)) return (ax, ay), (bx, by)