# Balabit 预训练 + Fine-tune 重构 Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Use Balabit Mouse Dynamics Challenge dataset to pretrain `TrajectoryFlowModel`, then fine-tune on user's 605 traces. Fix high-frequency lateral jitter and template-y Δt curves by aggressively removing deterministic post-processing. Add a quantitative eval module that produces Markdown reports with kinematic metrics and FFT spectra.

**Architecture:** Two-stage training (Balabit pretrain → 605 fine-tune via `resume_from`). New modules `ai_mouse/data_adapters/balabit.py` and `ai_mouse/eval/`. Existing model architecture (`TrajectoryFlowModel`), scroll subsystem, and frontend are unchanged. Data format remains identical to current `traces.jsonl`.

**Tech Stack:** Python 3.12+, PyTorch (CPU; optional CUDA), NumPy, SciPy, matplotlib (NEW dep for eval plots), uv package manager.

**Spec reference:** [docs/superpowers/specs/2026-05-10-balabit-pretrain-refactor-design.md](../specs/2026-05-10-balabit-pretrain-refactor-design.md)

**Prerequisites:**
1. User must download Balabit Mouse Dynamics Challenge dataset from https://github.com/balabit/Mouse-Dynamics-Challenge (clone or zip download). Plan assumes path is configurable via CLI arg.
2. Project will be initialized as git repo in Task 1 (currently not a git repo).

---

## Task 1: Project Setup (git init, dependencies, config)

**Files:**
- Create: `.gitignore`
- Modify: `pyproject.toml`
- Modify: `ai_mouse/config.py`
- Modify: `tests/conftest.py` (no functional change yet — just to verify import works after config changes)

- [ ] **Step 1: Initialize git repo (requires user OK)**

This step touches user-controlled state. Confirm with user before running.

```bash
cd /d/code/python/side/ai_mouse
git init
```

Expected: `Initialized empty Git repository in .../ai_mouse/.git/`

- [ ] **Step 2: Create .gitignore**

Create `.gitignore` at project root:

```
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
.venv/
.pytest_cache/

# IDE
.idea/
.vscode/

# uv
uv.lock.bak

# Data & models — large binary, do NOT commit
data/traces.jsonl
data/scroll_traces.jsonl
data/pretrain_traces.jsonl
data/models_v2/
data/models_v2_pretrained/
data/scroll_models/
data/eval_reports/
data/balabit_raw/

# OS
.DS_Store
Thumbs.db

# Playwright artifacts
.playwright-mcp/
```

- [ ] **Step 3: Add matplotlib to pyproject.toml**

Modify `pyproject.toml` — add matplotlib to `dependencies`:

```toml
[project]
name = "ai-mouse"
version = "0.1.0"
requires-python = ">=3.12,<3.14"
dependencies = [
    "torch>=2.2.0",
    "numpy>=1.26.0",
    "fastapi>=0.111.0",
    "uvicorn>=0.29.0",
    "scipy>=1.10.0",
    "matplotlib>=3.8.0",
]

[dependency-groups]
dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0", "httpx>=0.27.0"]
```

Then sync:

```bash
uv sync
```

Expected: `Resolved N packages` and matplotlib appears in `.venv`.

- [ ] **Step 4: Add new config dataclasses**

Append to `ai_mouse/config.py` (after existing `ServerConfig`):

```python
# ---------------------------------------------------------------------------
# Pretraining (Balabit) configuration
# ---------------------------------------------------------------------------


@dataclass
class PretrainConfig:
    """Hyperparameters for Balabit pretraining stage."""

    epochs: int = 200
    batch_size: int = 128
    lr: float = 3e-4
    seq_len: int = 64


@dataclass
class FinetuneConfig:
    """Hyperparameters for fine-tuning on user-collected data."""

    epochs: int = 50
    batch_size: int = 64
    lr: float = 1e-5  # 比预训练小一个数量级，防止灾难性遗忘
    seq_len: int = 64


# ---------------------------------------------------------------------------
# Balabit adapter configuration
# ---------------------------------------------------------------------------


@dataclass
class BalabitAdapterConfig:
    """Settings for Balabit CSV → traces.jsonl conversion."""

    window_ms: int = 1200       # click 前回溯窗口
    min_dist: int = 50          # 最小起终点距离 (px)
    min_events: int = 5         # 最小 Move 事件数
    max_span_ms: int = 5000     # 最大段时间跨度 (ms)
    max_gap_ms: int = 200       # 段内相邻 Move 最大时间差


# ---------------------------------------------------------------------------
# Eval configuration
# ---------------------------------------------------------------------------


@dataclass
class EvalConfig:
    """Settings for evaluation report generation."""

    n_samples: int = 1000
    fft_freq_band: tuple[float, float] = (4.0, 12.0)  # 生理震颤频段 (Hz)
    kl_bins: int = 50
```

- [ ] **Step 5: Verify imports**

```bash
uv run python -c "from ai_mouse.config import PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig; print('OK')"
```

Expected: `OK`

- [ ] **Step 6: Run all existing tests to confirm no regressions**

```bash
uv run pytest -x
```

Expected: all tests pass.

- [ ] **Step 7: Initial commit**

```bash
git add .gitignore pyproject.toml uv.lock ai_mouse/ tests/ static/ main.py docs/
git commit -m "chore: initialize git repo, add matplotlib dep, extend config

- Add .gitignore for Python/data/models
- Add matplotlib>=3.8.0 for eval plots
- Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses"
```

Expected: commit succeeds with hash printed.

---

## Task 2: Balabit Adapter — Core Module Skeleton + Types

**Files:**
- Create: `ai_mouse/data_adapters/__init__.py`
- Create: `ai_mouse/data_adapters/balabit.py`
- Create: `tests/test_balabit_adapter.py`

- [ ] **Step 1: Write failing tests for the public API surface**

Create `tests/test_balabit_adapter.py`:

```python
"""Tests for Balabit Mouse Dynamics Challenge data adapter."""
from __future__ import annotations

from pathlib import Path

import pytest


def test_module_exports():
    """The adapter module must export the public functions used by CLI."""
    from ai_mouse.data_adapters import balabit
    assert hasattr(balabit, "parse_session_csv")
    assert hasattr(balabit, "segment_by_clicks")
    assert hasattr(balabit, "filter_segments")
    assert hasattr(balabit, "process_session")
    assert hasattr(balabit, "MouseEvent")
    assert hasattr(balabit, "Segment")


def test_mouse_event_dataclass():
    """MouseEvent has expected fields."""
    from ai_mouse.data_adapters.balabit import MouseEvent
    e = MouseEvent(t_ms=100, button="NoButton", state="Move", x=300, y=400)
    assert e.t_ms == 100
    assert e.state == "Move"
    assert e.x == 300


def test_segment_dataclass():
    """Segment has expected fields."""
    from ai_mouse.data_adapters.balabit import MouseEvent, Segment
    events = [MouseEvent(t_ms=0, button="NoButton", state="Move", x=10, y=20)]
    s = Segment(events=events, click_x=100, click_y=200, click_t_ms=500, session_id="user1_s1")
    assert s.events == events
    assert s.click_x == 100
    assert s.session_id == "user1_s1"
```

- [ ] **Step 2: Run test, verify it fails**

```bash
uv run pytest tests/test_balabit_adapter.py -v
```

Expected: FAIL with `ModuleNotFoundError: No module named 'ai_mouse.data_adapters'`

- [ ] **Step 3: Create package skeleton**

Create `ai_mouse/data_adapters/__init__.py`:

```python
"""Data adapters: convert external datasets to the project's traces.jsonl format."""
```

Create `ai_mouse/data_adapters/balabit.py`:

```python
"""Adapter for the Balabit Mouse Dynamics Challenge dataset.

Source: https://github.com/balabit/Mouse-Dynamics-Challenge

Each session is a CSV file with columns:
    record timestamp, client timestamp, button, state, x, y

Where:
    state ∈ {Move, Pressed, Released, Drag, Scroll}
    button ∈ {NoButton, Left, Right, Wheel}

We extract "click-anchored" trajectory segments: each Pressed event
defines a target, and the W ms of Move events preceding it form one
training trace.
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class MouseEvent:
    """A single mouse event from a Balabit CSV row."""

    t_ms: int        # client timestamp in milliseconds (relative to session start)
    button: str      # "NoButton", "Left", "Right", "Wheel"
    state: str       # "Move", "Pressed", "Released", "Drag", "Scroll"
    x: int
    y: int


@dataclass
class Segment:
    """A click-anchored trajectory segment ready to be written to JSONL."""

    events: list[MouseEvent]   # only Move events, sorted by t_ms ascending
    click_x: int               # the Pressed event's x coordinate
    click_y: int               # the Pressed event's y coordinate
    click_t_ms: int            # the Pressed event's timestamp
    session_id: str            # e.g. "user7_session_42"


def parse_session_csv(path: Path) -> list[MouseEvent]:
    """Stub — implemented in Task 3."""
    raise NotImplementedError


def segment_by_clicks(
    events: list[MouseEvent],
    window_ms: int,
    session_id: str,
) -> list[Segment]:
    """Stub — implemented in Task 4."""
    raise NotImplementedError


def filter_segments(
    segments: list[Segment],
    min_events: int,
    min_dist: int,
    max_span_ms: int,
    max_gap_ms: int,
) -> list[Segment]:
    """Stub — implemented in Task 5."""
    raise NotImplementedError


def process_session(
    csv_path: Path,
    output_jsonl: Path,
    config,
) -> int:
    """Stub — implemented in Task 6."""
    raise NotImplementedError
```

- [ ] **Step 4: Run tests, verify they pass**

```bash
uv run pytest tests/test_balabit_adapter.py -v
```

Expected: 3 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/data_adapters/ tests/test_balabit_adapter.py
git commit -m "feat(adapter): scaffold balabit data adapter package"
```

---

## Task 3: Balabit Adapter — CSV Parsing

**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `parse_session_csv`)
- Modify: `tests/test_balabit_adapter.py` (add tests)

- [ ] **Step 1: Write failing tests for CSV parsing**

Append to `tests/test_balabit_adapter.py`:

```python
def _write_csv(path: Path, rows: list[str]) -> None:
    """Helper to write a Balabit-format CSV with header."""
    header = "record timestamp,client timestamp,button,state,x,y"
    path.write_text(header + "\n" + "\n".join(rows) + "\n", encoding="utf-8")


class TestParseSessionCsv:
    def test_parses_basic_rows(self, tmp_path):
        from ai_mouse.data_adapters.balabit import parse_session_csv
        csv = tmp_path / "session_1"
        _write_csv(csv, [
            "1500000000.000,0.000,NoButton,Move,100,200",
            "1500000000.050,0.050,NoButton,Move,110,210",
            "1500000000.100,0.100,Left,Pressed,120,220",
        ])
        events = parse_session_csv(csv)
        assert len(events) == 3
        assert events[0].t_ms == 0
        assert events[0].state == "Move"
        assert events[0].x == 100
        assert events[2].t_ms == 100
        assert events[2].state == "Pressed"
        assert events[2].button == "Left"

    def test_handles_float_timestamps(self, tmp_path):
        """Client timestamps are floats in seconds; we convert to int ms."""
        from ai_mouse.data_adapters.balabit import parse_session_csv
        csv = tmp_path / "session_2"
        _write_csv(csv, [
            "0,1.234,NoButton,Move,50,60",
            "0,1.250,NoButton,Move,55,65",
        ])
        events = parse_session_csv(csv)
        assert events[0].t_ms == 1234
        assert events[1].t_ms == 1250

    def test_skips_malformed_rows(self, tmp_path):
        """Rows with bad data are logged and skipped, not raised."""
        from ai_mouse.data_adapters.balabit import parse_session_csv
        csv = tmp_path / "session_3"
        _write_csv(csv, [
            "0,0.000,NoButton,Move,100,200",
            "BROKEN_ROW",
            "0,abc,NoButton,Move,100,200",   # bad timestamp
            "0,0.100,NoButton,Move,150,250",
        ])
        events = parse_session_csv(csv)
        assert len(events) == 2
        assert events[0].x == 100
        assert events[1].x == 150

    def test_returns_empty_list_for_empty_file(self, tmp_path):
        from ai_mouse.data_adapters.balabit import parse_session_csv
        csv = tmp_path / "session_4"
        csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8")
        events = parse_session_csv(csv)
        assert events == []
```

- [ ] **Step 2: Run tests, verify they fail**

```bash
uv run pytest tests/test_balabit_adapter.py::TestParseSessionCsv -v
```

Expected: FAIL with `NotImplementedError`.

- [ ] **Step 3: Implement parse_session_csv**

Replace the stub in `ai_mouse/data_adapters/balabit.py`:

```python
def parse_session_csv(path: Path) -> list[MouseEvent]:
    """Parse a Balabit session CSV file into MouseEvent objects.

    Malformed rows are logged and skipped (not raised).
    Client timestamps (seconds, float) are converted to int milliseconds.

    Args:
        path: Path to a Balabit session CSV file.

    Returns:
        List of MouseEvent in original order. Empty list if file is empty.
    """
    import csv as csv_module

    events: list[MouseEvent] = []
    with path.open("r", encoding="utf-8", newline="") as f:
        reader = csv_module.DictReader(f)
        for row_idx, row in enumerate(reader, 2):  # 1-based, header is line 1
            try:
                client_ts = float(row["client timestamp"])
                t_ms = int(round(client_ts * 1000))
                button = row["button"].strip()
                state = row["state"].strip()
                x = int(row["x"])
                y = int(row["y"])
            except (KeyError, ValueError, TypeError) as exc:
                logger.debug("Skipping malformed row %d in %s: %s", row_idx, path.name, exc)
                continue
            events.append(MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y))
    return events
```

- [ ] **Step 4: Run tests, verify they pass**

```bash
uv run pytest tests/test_balabit_adapter.py::TestParseSessionCsv -v
```

Expected: 4 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement Balabit CSV parser"
```

---

## Task 4: Balabit Adapter — Click-Anchored Segmentation

**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `segment_by_clicks`)
- Modify: `tests/test_balabit_adapter.py` (add tests)

- [ ] **Step 1: Write failing tests**

Append to `tests/test_balabit_adapter.py`:

```python
class TestSegmentByClicks:
    def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"):
        from ai_mouse.data_adapters.balabit import MouseEvent
        return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y)

    def test_one_click_one_segment(self):
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),
            self._make_event(100, "Move", 50, 60),
            self._make_event(500, "Move", 100, 100),
            self._make_event(600, "Pressed", 110, 110, button="Left"),
        ]
        segments = segment_by_clicks(events, window_ms=1200, session_id="test_s1")
        assert len(segments) == 1
        seg = segments[0]
        assert seg.click_x == 110
        assert seg.click_y == 110
        assert seg.click_t_ms == 600
        assert len(seg.events) == 3
        assert seg.session_id == "test_s1"

    def test_window_excludes_old_events(self):
        """Move events earlier than (click_t - window_ms) are dropped."""
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),       # too old
            self._make_event(100, "Move", 20, 30),     # too old
            self._make_event(900, "Move", 30, 40),     # in window
            self._make_event(1000, "Pressed", 40, 50, button="Left"),
        ]
        segments = segment_by_clicks(events, window_ms=200, session_id="s")
        assert len(segments) == 1
        assert len(segments[0].events) == 1
        assert segments[0].events[0].t_ms == 900

    def test_multiple_clicks_multiple_segments(self):
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),
            self._make_event(100, "Pressed", 50, 50, button="Left"),
            self._make_event(200, "Released", 50, 50, button="Left"),
            self._make_event(300, "Move", 60, 60),
            self._make_event(400, "Move", 70, 70),
            self._make_event(500, "Pressed", 80, 80, button="Left"),
        ]
        segments = segment_by_clicks(events, window_ms=1200, session_id="s")
        assert len(segments) == 2
        assert segments[0].click_x == 50
        assert segments[1].click_x == 80
        # Second segment's events must not include the first Pressed
        for e in segments[1].events:
            assert e.state == "Move"

    def test_skips_pressed_with_non_left_button(self):
        """Right-clicks and wheel-clicks don't anchor segments (only Left)."""
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),
            self._make_event(100, "Pressed", 50, 50, button="Right"),  # ignored
            self._make_event(200, "Move", 60, 60),
            self._make_event(300, "Pressed", 70, 70, button="Left"),   # anchor
        ]
        segments = segment_by_clicks(events, window_ms=1200, session_id="s")
        assert len(segments) == 1
        assert segments[0].click_x == 70

    def test_no_clicks_returns_empty(self):
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),
            self._make_event(100, "Move", 20, 30),
        ]
        segments = segment_by_clicks(events, window_ms=1200, session_id="s")
        assert segments == []

    def test_excludes_drag_events(self):
        """Drag events are not Move; segment should only include Move."""
        from ai_mouse.data_adapters.balabit import segment_by_clicks
        events = [
            self._make_event(0, "Move", 10, 20),
            self._make_event(100, "Drag", 30, 40),       # not Move
            self._make_event(200, "Move", 50, 60),
            self._make_event(300, "Pressed", 70, 80, button="Left"),
        ]
        segments = segment_by_clicks(events, window_ms=1200, session_id="s")
        assert len(segments) == 1
        # 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
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_balabit_adapter.py::TestSegmentByClicks -v
```

Expected: FAIL with `NotImplementedError`.

- [ ] **Step 3: Implement segment_by_clicks**

Replace the stub in `ai_mouse/data_adapters/balabit.py`:

```python
def segment_by_clicks(
    events: list[MouseEvent],
    window_ms: int,
    session_id: str,
) -> list[Segment]:
    """Extract click-anchored segments from a session.

    For each Left-button Pressed event, collect all Move events within
    [click_t - window_ms, click_t) into one segment.

    Args:
        events: Full session events (any state, any order is OK but typically sorted).
        window_ms: How far back to look before each click.
        session_id: String tag attached to every segment for debugging.

    Returns:
        List of Segment, one per Left Pressed event that has at least one preceding Move.
    """
    segments: list[Segment] = []
    for ev in events:
        if ev.state != "Pressed" or ev.button != "Left":
            continue
        click_t = ev.t_ms
        window_start = click_t - window_ms
        moves = [
            m for m in events
            if m.state == "Move" and window_start <= m.t_ms < click_t
        ]
        if not moves:
            continue
        moves.sort(key=lambda m: m.t_ms)
        segments.append(Segment(
            events=moves,
            click_x=ev.x,
            click_y=ev.y,
            click_t_ms=click_t,
            session_id=session_id,
        ))
    return segments
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_balabit_adapter.py::TestSegmentByClicks -v
```

Expected: 6 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement click-anchored segmentation"
```

---

## Task 5: Balabit Adapter — Filter Rules

**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `filter_segments`)
- Modify: `tests/test_balabit_adapter.py`

- [ ] **Step 1: Write failing tests**

Append to `tests/test_balabit_adapter.py`:

```python
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
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_balabit_adapter.py::TestFilterSegments -v
```

Expected: FAIL with `NotImplementedError`.

- [ ] **Step 3: Implement filter_segments**

Replace the stub in `ai_mouse/data_adapters/balabit.py`:

```python
def filter_segments(
    segments: list[Segment],
    min_events: int,
    min_dist: int,
    max_span_ms: int,
    max_gap_ms: int,
    coord_max: int = 5000,
) -> list[Segment]:
    """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
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_balabit_adapter.py::TestFilterSegments -v
```

Expected: 7 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement segment quality filters"
```

---

## Task 6: Balabit Adapter — Process Session + CLI

**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `process_session`, add `main`)
- Create: `ai_mouse/data_adapters/__main__.py`
- Modify: `tests/test_balabit_adapter.py`

- [ ] **Step 1: Write failing tests for process_session**

Append to `tests/test_balabit_adapter.py`:

```python
class TestProcessSession:
    def test_writes_jsonl_in_expected_format(self, tmp_path):
        from ai_mouse.config import BalabitAdapterConfig
        from ai_mouse.data_adapters.balabit import process_session

        # Construct a Balabit-format CSV with one valid segment
        csv_path = tmp_path / "user_session_42"
        rows = []
        # 10 Move events going from (100,100) to (500,200) over 500ms
        for i in range(10):
            t = i * 0.05
            x = 100 + i * 40
            y = 100 + i * 10
            rows.append(f"0,{t:.3f},NoButton,Move,{x},{y}")
        rows.append(f"0,0.550,Left,Pressed,510,210")
        _write_csv(csv_path, rows)

        out = tmp_path / "out.jsonl"
        config = BalabitAdapterConfig(
            window_ms=1200, min_dist=50, min_events=5,
            max_span_ms=5000, max_gap_ms=200,
        )
        n = process_session(csv_path, out, config)
        assert n == 1
        assert out.exists()

        line = out.read_text(encoding="utf-8").strip()
        record = json.loads(line)
        assert record["meta"]["start"] == [100, 100]
        assert record["meta"]["end"] == [510, 210]
        assert record["meta"]["dist"] > 0
        assert record["meta"]["source"] == "balabit"
        assert record["meta"]["session_id"] == "user_session_42"
        # Events: only move events, no down/up
        assert all(e["type"] == "move" for e in record["events"])
        # Timestamps relative (start at 0)
        assert record["events"][0]["t"] == 0

    def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
        from ai_mouse.config import BalabitAdapterConfig
        from ai_mouse.data_adapters.balabit import process_session

        csv_path = tmp_path / "empty_session"
        _write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"])  # no clicks

        out = tmp_path / "out.jsonl"
        config = BalabitAdapterConfig()
        n = process_session(csv_path, out, config)
        assert n == 0
        # File may not exist or may be empty — both acceptable
        if out.exists():
            assert out.read_text() == ""

    def test_appends_to_existing_jsonl(self, tmp_path):
        from ai_mouse.config import BalabitAdapterConfig
        from ai_mouse.data_adapters.balabit import process_session

        out = tmp_path / "out.jsonl"
        out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")

        csv_path = tmp_path / "session_x"
        rows = [f"0,{i*0.05:.3f},NoButton,Move,{100+i*40},100" for i in range(10)]
        rows.append("0,0.550,Left,Pressed,510,100")
        _write_csv(csv_path, rows)

        config = BalabitAdapterConfig()
        process_session(csv_path, out, config)

        lines = out.read_text(encoding="utf-8").strip().split("\n")
        assert len(lines) == 2  # original + appended
```

`json` import — make sure top of test file already has `import json`. If not, add.

Add this near the top of `tests/test_balabit_adapter.py`:

```python
import json
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_balabit_adapter.py::TestProcessSession -v
```

Expected: FAIL with `NotImplementedError`.

- [ ] **Step 3: Implement process_session**

Replace the `process_session` stub in `ai_mouse/data_adapters/balabit.py`:

```python
def process_session(
    csv_path: Path,
    output_jsonl: Path,
    config,
) -> int:
    """Convert one Balabit session CSV to JSONL traces, append to output.

    Args:
        csv_path: Path to a Balabit session CSV.
        output_jsonl: Output JSONL file (will be appended to).
        config: BalabitAdapterConfig with window_ms / min_dist / etc.

    Returns:
        Number of valid segments written.
    """
    import math

    session_id = csv_path.stem  # e.g. "session_42"
    events = parse_session_csv(csv_path)
    if not events:
        return 0

    raw_segments = segment_by_clicks(
        events, window_ms=config.window_ms, session_id=session_id
    )
    valid_segments = filter_segments(
        raw_segments,
        min_events=config.min_events,
        min_dist=config.min_dist,
        max_span_ms=config.max_span_ms,
        max_gap_ms=config.max_gap_ms,
    )

    if not valid_segments:
        return 0

    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
    n_written = 0
    with output_jsonl.open("a", encoding="utf-8") as f:
        for seg in valid_segments:
            sx, sy = seg.events[0].x, seg.events[0].y
            ex, ey = seg.click_x, seg.click_y
            dist = math.hypot(ex - sx, ey - sy)
            angle = math.degrees(math.atan2(ey - sy, ex - sx))
            t0 = seg.events[0].t_ms
            record = {
                "meta": {
                    "start": [sx, sy],
                    "end": [ex, ey],
                    "dist": int(round(dist)),
                    "angle": round(angle, 1),
                    "source": "balabit",
                    "session_id": seg.session_id,
                },
                "events": [
                    {"type": "move", "x": e.x, "y": e.y, "t": e.t_ms - t0}
                    for e in seg.events
                ],
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
            n_written += 1
    return n_written
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_balabit_adapter.py::TestProcessSession -v
```

Expected: 3 tests pass.

- [ ] **Step 5: Add CLI main function to balabit.py**

Append to `ai_mouse/data_adapters/balabit.py`:

```python
def main(argv: list[str] | None = None) -> int:
    """CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
    import argparse

    from ai_mouse.config import BalabitAdapterConfig

    parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
    parser.add_argument(
        "--input", type=Path, required=True,
        help="Directory containing Balabit session CSV files (recursive)",
    )
    parser.add_argument(
        "--output", type=Path, default=Path("data/pretrain_traces.jsonl"),
        help="Output JSONL path (default: data/pretrain_traces.jsonl)",
    )
    parser.add_argument("--window-ms", type=int, default=1200)
    parser.add_argument("--min-dist", type=int, default=50)
    parser.add_argument("--min-events", type=int, default=5)
    parser.add_argument("--max-span-ms", type=int, default=5000)
    parser.add_argument("--max-gap-ms", type=int, default=200)
    parser.add_argument(
        "--overwrite", action="store_true",
        help="Truncate output file before writing (default: append)",
    )
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    if not args.input.is_dir():
        logger.error("Input is not a directory: %s", args.input)
        return 2

    config = BalabitAdapterConfig(
        window_ms=args.window_ms,
        min_dist=args.min_dist,
        min_events=args.min_events,
        max_span_ms=args.max_span_ms,
        max_gap_ms=args.max_gap_ms,
    )

    if args.overwrite and args.output.exists():
        args.output.unlink()

    # Recursively walk the input directory looking for session files.
    # Balabit files have no extension; we accept any regular file.
    csv_files = sorted([p for p in args.input.rglob("*") if p.is_file() and not p.name.startswith(".")])
    if not csv_files:
        logger.error("No session files found under %s", args.input)
        return 2

    total = 0
    for i, csv_path in enumerate(csv_files, 1):
        try:
            n = process_session(csv_path, args.output, config)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Skipping %s due to error: %s", csv_path.name, exc)
            continue
        total += n
        if i % 10 == 0 or i == len(csv_files):
            logger.info("Processed %d/%d sessions, %d segments so far", i, len(csv_files), total)

    logger.info("Done. Wrote %d segments to %s", total, args.output)
    return 0
```

- [ ] **Step 6: Create __main__.py for CLI dispatch**

Create `ai_mouse/data_adapters/__main__.py`:

```python
"""CLI dispatch: `python -m ai_mouse.data_adapters.balabit ...`

Note: This file makes `python -m ai_mouse.data_adapters` invokable but for
clarity prefer the explicit form `python -m ai_mouse.data_adapters.balabit`.
"""
from __future__ import annotations

import sys

from ai_mouse.data_adapters.balabit import main

if __name__ == "__main__":
    sys.exit(main())
```

Also add a `__main__` guard at the bottom of `balabit.py` so `python -m ai_mouse.data_adapters.balabit` works directly:

```python
if __name__ == "__main__":
    import sys
    sys.exit(main())
```

- [ ] **Step 7: Smoke-test the CLI**

```bash
uv run python -m ai_mouse.data_adapters.balabit --help
```

Expected: argparse help text printed.

- [ ] **Step 8: Run all adapter tests**

```bash
uv run pytest tests/test_balabit_adapter.py -v
```

Expected: all tests pass.

- [ ] **Step 9: Commit**

```bash
git add ai_mouse/data_adapters/ tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement process_session and CLI"
```

---

## Task 7: Trainer — TrajectoryDataset (Streaming Augmentation)

**Files:**
- Modify: `ai_mouse/trainer.py` (replace `_augment` usage with `TrajectoryDataset`)
- Modify: `tests/test_trainer.py` (add tests, keep existing ones passing)

- [ ] **Step 1: Write failing test for the new Dataset class**

Append to `tests/test_trainer.py`:

```python
class TestTrajectoryDataset:
    def test_dataset_length_with_augmentation(self):
        """Dataset length = N * 6 when augment=True."""
        from ai_mouse.trainer import TrajectoryDataset
        seq = np.zeros((10, 64, 3), dtype=np.float32)
        cond = np.zeros((10, 3), dtype=np.float32)
        ds = TrajectoryDataset(seq, cond, augment=True)
        assert len(ds) == 60

    def test_dataset_length_without_augmentation(self):
        from ai_mouse.trainer import TrajectoryDataset
        seq = np.zeros((10, 64, 3), dtype=np.float32)
        cond = np.zeros((10, 3), dtype=np.float32)
        ds = TrajectoryDataset(seq, cond, augment=False)
        assert len(ds) == 10

    def test_getitem_returns_tensors(self):
        from ai_mouse.trainer import TrajectoryDataset
        import torch
        seq = np.random.randn(5, 64, 3).astype(np.float32)
        cond = np.random.randn(5, 3).astype(np.float32)
        ds = TrajectoryDataset(seq, cond, augment=True)
        s, c = ds[0]
        assert isinstance(s, torch.Tensor)
        assert isinstance(c, torch.Tensor)
        assert s.shape == (64, 3)
        assert c.shape == (3,)

    def test_aug_id_zero_returns_original(self):
        """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
        from ai_mouse.trainer import TrajectoryDataset
        import torch
        seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
        cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
        ds = TrajectoryDataset(seq, cond, augment=True)
        s0, c0 = ds[0]
        np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5)
        np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5)

    def test_aug_id_one_flips_lateral(self):
        """Aug id 1 should flip the sign of the lateral channel (index 1)."""
        from ai_mouse.trainer import TrajectoryDataset
        seq = np.zeros((1, 64, 3), dtype=np.float32)
        seq[0, :, 1] = 0.5  # lateral all positive
        cond = np.zeros((1, 3), dtype=np.float32)
        ds = TrajectoryDataset(seq, cond, augment=True)
        # idx=1 → base_idx=0, aug_id=1 → flip
        s1, _ = ds[1]
        assert (s1[:, 1] < 0).all()
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_trainer.py::TestTrajectoryDataset -v
```

Expected: FAIL with `ImportError` (TrajectoryDataset not defined).

- [ ] **Step 3: Implement TrajectoryDataset class**

Add to `ai_mouse/trainer.py`, just below the `_augment` function (keep `_augment` for now — Task removes it):

```python
class TrajectoryDataset(torch.utils.data.Dataset):
    """Trajectory dataset with on-the-fly 6× augmentation.

    Replaces the old eager `_augment(seq, cond)` which expanded the dataset
    6× in memory before training. With this class, the original (N, T, 3)
    arrays stay as-is and each `__getitem__` call computes one of the 6
    augmentation variants on demand.

    Augmentation variants (matching legacy `_augment` semantics):
      0 — original
      1 — lateral flip (lateral → −lateral)
      2 — speed ×0.8  (log_dt[1:] += log(1.25),  cond[2] += log(1.25))
      3 — speed ×1.2  (log_dt[1:] += log(1/1.2), cond[2] += log(1/1.2))
      4 — temporal noise (log_dt[1:] += N(0, 0.05))
      5 — flip + speed ×0.9 (lateral flip, log_dt[1:] += log(1/0.9), cond[2] += log(1/0.9))
    """

    _LOG_1_25 = math.log(1.25)
    _LOG_INV_1_2 = math.log(1.0 / 1.2)
    _LOG_1_1 = math.log(1.0 / 0.9)

    def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
        self.seq = seq
        self.cond = cond
        self.augment = augment
        self._n_aug = 6 if augment else 1

    def __len__(self) -> int:
        return len(self.seq) * self._n_aug

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        base = idx // self._n_aug
        aug_id = idx % self._n_aug

        s = self.seq[base].copy()
        c = self.cond[base].copy()

        if aug_id == 1:
            s[:, 1] = -s[:, 1]
        elif aug_id == 2:
            s[1:, 2] += self._LOG_1_25
            c[2] += self._LOG_1_25
        elif aug_id == 3:
            s[1:, 2] += self._LOG_INV_1_2
            c[2] += self._LOG_INV_1_2
        elif aug_id == 4:
            noise = np.random.normal(0.0, 0.05, size=s[1:, 2].shape).astype(np.float32)
            s[1:, 2] += noise
        elif aug_id == 5:
            s[:, 1] = -s[:, 1]
            s[1:, 2] += self._LOG_1_1
            c[2] += self._LOG_1_1

        return torch.from_numpy(s), torch.from_numpy(c)
```

- [ ] **Step 4: Run new tests, verify pass**

```bash
uv run pytest tests/test_trainer.py::TestTrajectoryDataset -v
```

Expected: 5 tests pass.

- [ ] **Step 5: Switch `train()` to use TrajectoryDataset**

In `ai_mouse/trainer.py`, find the existing block (around line 332):

```python
    # ---- Augment ----
    seq_np, cond_np = _augment(seq_np, cond_np)
    logger.info("After augmentation: %d samples", len(seq_np))

    seq_t = torch.from_numpy(seq_np)    # (N, seq_len, 3)
    cond_t = torch.from_numpy(cond_np)  # (N, 3)
```

Replace with:

```python
    # ---- Build streaming dataset (on-the-fly 6× augmentation) ----
    if config.augment:
        logger.info("Using on-the-fly 6× augmentation, base samples: %d", len(seq_np))
    ds = TrajectoryDataset(seq_np, cond_np, augment=config.augment)
    logger.info("Effective dataset size: %d", len(ds))
```

Then find the existing DataLoader construction (around line 353):

```python
    ds = TensorDataset(seq_t, cond_t)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
```

Replace with:

```python
    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
```

(The new `ds` already exists from the previous edit — just remove the TensorDataset construction.)

Also remove the now-unused import line `from torch.utils.data import DataLoader, TensorDataset` and replace with:

```python
from torch.utils.data import DataLoader
```

- [ ] **Step 6: Verify legacy augmentation tests still pass**

The existing test `TestAugment::test_augmentation_multiplies_data` calls `_augment` directly, which we kept. Run all trainer tests:

```bash
uv run pytest tests/test_trainer.py -v
```

Expected: all tests pass (including TestAugment, TestLoadAndPrepare, TestTrain, TestTrajectoryDataset).

- [ ] **Step 7: Commit**

```bash
git add ai_mouse/trainer.py tests/test_trainer.py
git commit -m "feat(trainer): replace eager _augment with streaming TrajectoryDataset"
```

---

## Task 8: Trainer — Resume from Checkpoint

**Files:**
- Modify: `ai_mouse/trainer.py` (add `resume_from` parameter to `train()`)
- Modify: `tests/test_trainer.py`

- [ ] **Step 1: Write failing test for resume_from**

Append to `tests/test_trainer.py`:

```python
class TestResumeFrom:
    def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):
        """train() with resume_from should load weights from given checkpoint dir."""
        import torch
        from ai_mouse.trainer import train
        from ai_mouse.models import TrajectoryFlowModel

        # First, train an initial model and save it
        ckpt_dir = tmp_path / "pretrain"
        train(
            data_path=synthetic_traces_file,
            output_dir=ckpt_dir,
            epochs=2,
            batch_size=8,
            seq_len=64,
        )
        assert (ckpt_dir / "flow_model.pt").exists()

        # Read its weights to compare later
        m_pretrain = TrajectoryFlowModel(seq_len=64)
        m_pretrain.load_state_dict(torch.load(ckpt_dir / "flow_model.pt", weights_only=True))
        first_param_pre = next(m_pretrain.parameters()).clone()

        # Now train with resume_from for 0 epochs — weights should still be loaded
        out_dir = tmp_path / "finetune"
        train(
            data_path=synthetic_traces_file,
            output_dir=out_dir,
            epochs=1,
            batch_size=8,
            seq_len=64,
            resume_from=ckpt_dir,
        )

        m_after = TrajectoryFlowModel(seq_len=64)
        m_after.load_state_dict(torch.load(out_dir / "flow_model.pt", weights_only=True))
        first_param_after = next(m_after.parameters())

        # After 1 epoch, weights should be close to pre-train, not random init
        # (random init would be O(1) magnitude apart; 1 epoch on small data shifts O(0.1))
        diff = (first_param_pre - first_param_after).abs().mean().item()
        assert diff < 0.5, f"Resume_from weights diverged too much: {diff}"

    def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path):
        from ai_mouse.trainer import train

        with pytest.raises(FileNotFoundError):
            train(
                data_path=synthetic_traces_file,
                output_dir=tmp_path / "out",
                epochs=1,
                batch_size=8,
                seq_len=64,
                resume_from=tmp_path / "nonexistent",
            )
```

- [ ] **Step 2: Run test, verify failure**

```bash
uv run pytest tests/test_trainer.py::TestResumeFrom -v
```

Expected: FAIL with `TypeError: train() got an unexpected keyword argument 'resume_from'`.

- [ ] **Step 3: Add resume_from parameter to train()**

In `ai_mouse/trainer.py`, modify the `train` function signature:

```python
def train(
    data_path: Path,
    output_dir: Path,
    epochs: int = 300,
    batch_size: int = 64,
    lr: float = 3e-4,
    seq_len: int = 64,
    progress_callback: Callable[[dict], None] | None = None,
    config: TrainConfig | None = None,
    resume_from: Path | None = None,
) -> None:
```

Update the docstring's Args section to add:

```
        resume_from:        if given, load model weights from this checkpoint
                            directory (must contain flow_model.pt). Used for
                            two-stage training (pretrain → fine-tune).
```

Then, after model construction (find the line `model = TrajectoryFlowModel(...)` around line 341) and BEFORE the optimiser is created, insert:

```python
    # ---- Resume from checkpoint if requested ----
    if resume_from is not None:
        resume_path = Path(resume_from) / "flow_model.pt"
        if not resume_path.exists():
            raise FileNotFoundError(
                f"resume_from checkpoint not found: {resume_path}"
            )
        logger.info("Resuming from checkpoint: %s", resume_path)
        state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict)
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_trainer.py::TestResumeFrom -v
```

Expected: 2 tests pass.

- [ ] **Step 5: Run all trainer tests**

```bash
uv run pytest tests/test_trainer.py -v
```

Expected: all tests pass.

- [ ] **Step 6: Commit**

```bash
git add ai_mouse/trainer.py tests/test_trainer.py
git commit -m "feat(trainer): add resume_from for two-stage training"
```

---

## Task 9: Server — Auto-Resume When Pretrained Checkpoint Exists

**Files:**
- Modify: `ai_mouse/server/routes_train.py`
- Modify: `tests/test_server.py` (verify still passes)

- [ ] **Step 1: Update routes_train.py to detect and pass resume_from**

In `ai_mouse/server/routes_train.py`, modify `_paths()` to also return the pretrained dir:

Find:

```python
def _paths() -> tuple[Path, Path]:
    data_dir = get_data_dir()
    return data_dir / "traces.jsonl", data_dir / "models_v2"
```

Replace with:

```python
def _paths() -> tuple[Path, Path, Path]:
    data_dir = get_data_dir()
    return (
        data_dir / "traces.jsonl",
        data_dir / "models_v2",
        data_dir / "models_v2_pretrained",
    )
```

Update `_trace_count()` and `_model_trained()`:

```python
def _trace_count() -> int:
    traces_path, _, _ = _paths()
    if not traces_path.exists():
        return 0
    return sum(
        1
        for line in traces_path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    )


def _model_trained() -> bool:
    _, models_dir, _ = _paths()
    return (models_dir / "flow_model.pt").exists()
```

Then update `_train_sse_generator`'s inner `run_training_async`:

Find:

```python
    async def run_training_async() -> None:
        from ai_mouse.trainer import train

        traces_path, models_dir = _paths()
        data_path = Path(req.data_path) if req.data_path else traces_path
        output_dir = Path(req.output_dir) if req.output_dir else models_dir
        try:
            await asyncio.to_thread(
                train,
                data_path=data_path,
                output_dir=output_dir,
                epochs=req.epochs,
                progress_callback=callback,
            )
        except Exception as exc:  # noqa: BLE001
            queue.put_nowait({"error": str(exc)})
```

Replace with:

```python
    async def run_training_async() -> None:
        from ai_mouse.trainer import train

        traces_path, models_dir, pretrained_dir = _paths()
        data_path = Path(req.data_path) if req.data_path else traces_path
        output_dir = Path(req.output_dir) if req.output_dir else models_dir

        # Auto-detect pretrained checkpoint and switch to fine-tune mode
        resume_from: Path | None = None
        effective_lr = 3e-4
        if (pretrained_dir / "flow_model.pt").exists():
            resume_from = pretrained_dir
            effective_lr = 1e-5  # fine-tune lr
            logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
            queue.put_nowait({
                "info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
                        f"running fine-tune at lr={effective_lr}",
            })

        try:
            await asyncio.to_thread(
                train,
                data_path=data_path,
                output_dir=output_dir,
                epochs=req.epochs,
                lr=effective_lr,
                progress_callback=callback,
                resume_from=resume_from,
            )
        except Exception as exc:  # noqa: BLE001
            queue.put_nowait({"error": str(exc)})
```

- [ ] **Step 2: Run server tests, verify still pass**

```bash
uv run pytest tests/test_server.py -v
```

Expected: all tests pass. The auto-resume behaviour is exercised when `models_v2_pretrained` exists, which is not the case in the default test environment, so the existing tests still hit the from-scratch path.

- [ ] **Step 3: Run full test suite**

```bash
uv run pytest -x
```

Expected: all tests pass.

- [ ] **Step 4: Commit**

```bash
git add ai_mouse/server/routes_train.py
git commit -m "feat(server): auto-resume from pretrained checkpoint when available"
```

---

## Task 10: Generator — Remove speed_profile and median±1.1 Hard Clip

**Files:**
- Modify: `ai_mouse/generator.py`
- Modify: `tests/test_generator.py`

- [ ] **Step 1: Update existing tests to be tolerant of the new behaviour**

The current tests use freshly-initialised (untrained) weights and only check structural properties (timestamps monotonic, endpoints close, click events present). These should keep passing — verify.

```bash
uv run pytest tests/test_generator.py -v
```

Expected: all tests pass.

Add a new test that explicitly checks the OLD speed_profile is GONE (regression guard):

Append to `tests/test_generator.py`:

```python
class TestPostProcessing:
    def test_dt_diversity_preserved(self, model_dir):
        """After removing speed_profile + median clip, multiple generations
        should differ in their Δt sequences (not all identical)."""
        results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
                   for _ in range(5)]
        # Extract Δt sequences (only move events, not click events)
        dts = []
        for r in results:
            moves = r[:-2]
            dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)]
            dts.append(dt_seq)
        # At least 2 of the 5 sequences should differ at any given index
        # (untrained model produces noisy outputs but post-processing must
        # not collapse them to the same template)
        for i in range(min(len(d) for d in dts)):
            values = {tuple([d[i]]) for d in dts}
            if len(values) > 1:
                return  # at least one position has variation — pass
        pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")
```

- [ ] **Step 2: Run test, may pass for untrained model but verifies regression guard**

```bash
uv run pytest tests/test_generator.py::TestPostProcessing -v
```

Expected: PASS (untrained models produce random outputs, so they vary).

- [ ] **Step 3: Remove the speed_profile and median clip blocks**

In `ai_mouse/generator.py`, find this block (lines ~261-272):

```python
    # The model tends to produce exaggerated deceleration at the tail
    # (last 10 points log_dt ~3-5 vs middle ~1.5).
    # Cap the max-to-median ratio to ~3× (i.e., tail Δt ≤ 3× median Δt)
    median_ldt = float(np.median(log_dt[1:]))
    # Allow max log_dt = median + 1.1 (exp(1.1) ≈ 3× ratio)
    max_allowed = median_ldt + 1.1
    min_allowed = max(median_ldt - 1.1, 0.0)
    for i in range(1, len(log_dt)):
        if log_dt[i] > max_allowed:
            log_dt[i] = max_allowed
        elif log_dt[i] < min_allowed:
            log_dt[i] = min_allowed
```

DELETE this entire block.

Then find the speed_profile block immediately after (lines ~273-286):

```python
    # Apply asymmetric speed profile: start slow, fast in middle, gentle end
    # Mimics natural mouse movement (accelerate → cruise → decelerate)
    t_frac = np.linspace(0, 1, len(log_dt))
    speed_profile = np.zeros_like(log_dt, dtype=float)
    for i in range(1, len(log_dt)):
        t = t_frac[i]
        if t < 0.15:
            # Acceleration phase: start slow (+0.3 at t=0, → 0 at t=0.15)
            speed_profile[i] = 0.3 * (1.0 - t / 0.15)
        elif t > 0.85:
            # Deceleration phase: end slightly slow (+0.2 at t=1)
            speed_profile[i] = 0.2 * ((t - 0.85) / 0.15)
        # Middle: speed_profile = 0 (fastest, no penalty)
    log_dt[1:] = log_dt[1:] + speed_profile[1:]
```

DELETE this entire block too.

- [ ] **Step 4: Update the docstring at the top of the file**

In `ai_mouse/generator.py`, find the module docstring (lines 1-22). Replace it with:

```python
"""Inference layer: Flow Matching trajectory generation.

Pipeline:
1. Load model from model_dir (flow_model.pt, click_dist.json,
   duration_dist.json, train_config.json).
2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)].
3. Sample total_duration from duration_dist.json by distance bin (log-normal).
4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory.
5. Spatial post-processing:
   a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
   b. Smooth start: dampen lateral near start (first 4 points).
   c. Enforce forward monotonicity (prevent x-axis jitter).
6. Temporal post-processing:
   a. Clip log_dt to [0, 5] to prevent exponential explosion.
   (speed profile and median±1.1 hard clip removed in 2026-05 refactor —
    let the model's learned timing distribution come through naturally.)
7. Decode to pixels via decode_trajectory.
8. Resample to n_points if n_points != model seq_len.
9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150].
10. Ensure timestamps monotonically increasing.
11. Append click events sampled from truncated normal.
"""
```

- [ ] **Step 5: Run tests**

```bash
uv run pytest tests/test_generator.py -v
```

Expected: all tests pass.

- [ ] **Step 6: Commit**

```bash
git add ai_mouse/generator.py tests/test_generator.py
git commit -m "refactor(generator): remove deterministic speed_profile and hard log_dt clip

These post-processing hacks were added to compensate for small-data
training. With Balabit pretraining they suppress the multimodal
timing distribution and cause the template-y Δt curves seen in the
verify UI."
```

---

## Task 11: Generator — Add Lateral Gaussian Smoothing

**Files:**
- Modify: `ai_mouse/generator.py`
- Modify: `tests/test_generator.py`

- [ ] **Step 1: Write failing test for the smoothing helper**

Append to `tests/test_generator.py`:

```python
class TestGaussianSmooth:
    def test_endpoints_preserved(self):
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        assert smoothed[0] == 1.0
        assert smoothed[-1] == 2.0

    def test_smooths_high_frequency(self):
        """A high-frequency square wave should have reduced amplitude after smoothing."""
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        # Interior amplitude should be reduced
        interior_orig = x[2:-2]
        interior_smooth = smoothed[2:-2]
        assert interior_smooth.std() < interior_orig.std()

    def test_constant_signal_unchanged(self):
        from ai_mouse.generator import _gaussian_smooth
        x = np.full(20, 0.5, dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        np.testing.assert_allclose(smoothed, x, rtol=1e-6)

    def test_short_array_returns_unchanged(self):
        """Arrays shorter than the kernel are returned unchanged."""
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        np.testing.assert_allclose(smoothed, x, rtol=1e-6)
```

- [ ] **Step 2: Run test, verify failure**

```bash
uv run pytest tests/test_generator.py::TestGaussianSmooth -v
```

Expected: FAIL with `ImportError: cannot import name '_gaussian_smooth'`.

- [ ] **Step 3: Implement _gaussian_smooth**

Add to `ai_mouse/generator.py`, just below the imports (before any function definitions):

```python
def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-point gaussian smoothing along a 1-D array, preserving endpoints.

    Args:
        x: 1-D input array.
        sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
               weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].

    Returns:
        Smoothed array of the same shape. x[0] and x[-1] are unchanged.
        If len(x) < 5, returns x unchanged (kernel won't fit).
    """
    if len(x) < 5:
        return x.copy()
    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
    kernel /= kernel.sum()
    smoothed = np.convolve(x, kernel, mode="same")
    smoothed[0] = x[0]
    smoothed[-1] = x[-1]
    return smoothed
```

- [ ] **Step 4: Run smoothing tests, verify pass**

```bash
uv run pytest tests/test_generator.py::TestGaussianSmooth -v
```

Expected: 4 tests pass.

- [ ] **Step 5: Apply smoothing to lateral in generate()**

In `ai_mouse/generator.py`, find the spatial post-processing block. After this code:

```python
    # Clamp forward to [0, 1] and re-force endpoints after monotonicity fix
    forward = np.clip(forward, 0.0, 1.0)
    forward[0] = 0.0
    forward[-1] = 1.0
```

Add:

```python
    # Lateral 5-point gaussian smoothing (endpoints preserved)
    lateral = _gaussian_smooth(lateral, sigma=1.0)
```

- [ ] **Step 6: Update module docstring to document the new smoothing step**

In the module docstring at the top of `ai_mouse/generator.py`, find:

```python
5. Spatial post-processing:
   a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
   b. Smooth start: dampen lateral near start (first 4 points).
   c. Enforce forward monotonicity (prevent x-axis jitter).
```

Replace with:

```python
5. Spatial post-processing:
   a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
   b. Smooth start: dampen lateral near start (first 4 points).
   c. Enforce forward monotonicity (prevent x-axis jitter).
   d. 5-point gaussian smooth on lateral (preserve endpoints).
```

- [ ] **Step 7: Run all generator tests**

```bash
uv run pytest tests/test_generator.py -v
```

Expected: all tests pass.

- [ ] **Step 8: Commit**

```bash
git add ai_mouse/generator.py tests/test_generator.py
git commit -m "feat(generator): add 5-point gaussian smoothing on lateral"
```

---

## Task 12: Eval — Kinematics Metrics

**Files:**
- Create: `ai_mouse/eval/__init__.py`
- Create: `ai_mouse/eval/metrics.py`
- Create: `tests/test_eval_metrics.py`

- [ ] **Step 1: Write failing tests for kinematics metrics**

Create `tests/test_eval_metrics.py`:

```python
"""Tests for the eval metrics module."""
from __future__ import annotations

import numpy as np
import pytest


class TestKinematics:
    def test_compute_speed_constant_velocity(self):
        """Constant-velocity trajectory has constant speed."""
        from ai_mouse.eval.metrics import compute_speed
        # 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms
        xs = np.arange(0, 100, 10, dtype=float)
        ys = np.zeros(10, dtype=float)
        ts = np.arange(0, 1000, 100, dtype=float)
        v = compute_speed(xs, ys, ts)
        # All speeds should be ≈ 0.1 px/ms
        assert v.shape == (9,)  # n-1 differences
        np.testing.assert_allclose(v, 0.1, rtol=1e-4)

    def test_compute_speed_handles_zero_dt(self):
        """Adjacent points with same timestamp must not produce NaN/inf."""
        from ai_mouse.eval.metrics import compute_speed
        xs = np.array([0.0, 10.0, 20.0])
        ys = np.array([0.0, 0.0, 0.0])
        ts = np.array([0.0, 0.0, 100.0])  # zero dt between [0] and [1]
        v = compute_speed(xs, ys, ts)
        assert np.isfinite(v).all()

    def test_compute_acceleration(self):
        """Linearly increasing speed → constant acceleration."""
        from ai_mouse.eval.metrics import compute_acceleration
        # speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms²
        speeds = np.array([0.1, 0.2, 0.3, 0.4])
        ts = np.array([100.0, 200.0, 300.0, 400.0])
        a = compute_acceleration(speeds, ts)
        np.testing.assert_allclose(a, 0.001, rtol=1e-4)

    def test_compute_jerk(self):
        from ai_mouse.eval.metrics import compute_jerk
        # accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001
        accels = np.array([0.001, 0.002, 0.003])
        ts = np.array([200.0, 300.0, 400.0])
        j = compute_jerk(accels, ts)
        np.testing.assert_allclose(j, 1e-5, rtol=1e-4)


class TestStatsSummary:
    def test_compute_stats_returns_expected_keys(self):
        from ai_mouse.eval.metrics import compute_stats
        x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        s = compute_stats(x)
        assert "mean" in s
        assert "std" in s
        assert "cv" in s
        assert "p25" in s
        assert "p50" in s
        assert "p75" in s
        assert "p95" in s

    def test_cv_for_constant_is_zero(self):
        from ai_mouse.eval.metrics import compute_stats
        x = np.full(10, 3.0)
        s = compute_stats(x)
        assert s["cv"] == 0.0
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_eval_metrics.py::TestKinematics tests/test_eval_metrics.py::TestStatsSummary -v
```

Expected: FAIL with `ModuleNotFoundError: No module named 'ai_mouse.eval'`.

- [ ] **Step 3: Create the eval package and metrics module**

Create `ai_mouse/eval/__init__.py`:

```python
"""Evaluation module: kinematic metrics and Markdown report generation."""
```

Create `ai_mouse/eval/metrics.py`:

```python
"""Kinematic metrics for mouse trajectory evaluation.

All inputs are 1-D NumPy arrays. Time is in milliseconds, position in pixels.
Velocities are px/ms, accelerations px/ms², jerks px/ms³.
"""
from __future__ import annotations

import numpy as np


def compute_speed(
    xs: np.ndarray, ys: np.ndarray, ts: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
    """Compute scalar speed at each step.

    Args:
        xs: (N,) x coordinates.
        ys: (N,) y coordinates.
        ts: (N,) timestamps in ms.
        eps: minimum dt (ms) to avoid div-by-zero.

    Returns:
        (N-1,) array of speeds (px/ms).
    """
    dx = np.diff(xs)
    dy = np.diff(ys)
    dt = np.maximum(np.diff(ts), eps)
    return np.hypot(dx, dy) / dt


def compute_acceleration(speeds: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Compute scalar acceleration from speeds.

    Args:
        speeds: (M,) speeds (px/ms). Typically M = N-1 from compute_speed.
        ts: (M+1,) timestamps that produced those speeds (ms).
            We use the midpoints between adjacent ts.
        eps: minimum dt (ms) to avoid div-by-zero.

    Returns:
        (M-1,) array of accelerations (px/ms²).
    """
    if len(speeds) < 2:
        return np.array([], dtype=float)
    # ts has length M+1; we need M-1 dts between speed points
    # speed[i] is between ts[i] and ts[i+1], so it's at ts midpoint (ts[i]+ts[i+1])/2
    midpoints = (ts[:-1] + ts[1:]) / 2.0
    dt = np.maximum(np.diff(midpoints), eps)
    return np.diff(speeds) / dt


def compute_jerk(accels: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Compute jerk from accelerations.

    Args:
        accels: (K,) accelerations.
        ts: (K+2,) timestamps that produced those accelerations.
            Used to derive midpoint-of-midpoint dts.
        eps: minimum dt to avoid div-by-zero.

    Returns:
        (K-1,) array of jerks (px/ms³).
    """
    if len(accels) < 2:
        return np.array([], dtype=float)
    # Approximate dt for jerks as average dt of original ts (good enough for stats)
    dt_avg = np.maximum(np.diff(ts).mean(), eps)
    return np.diff(accels) / dt_avg


def compute_stats(x: np.ndarray) -> dict[str, float]:
    """Summary statistics for a 1-D distribution.

    Returns:
        dict with keys: mean, std, cv (coef of variation), p25, p50, p75, p95.
    """
    if len(x) == 0:
        return {k: 0.0 for k in ("mean", "std", "cv", "p25", "p50", "p75", "p95")}
    x = np.asarray(x, dtype=float)
    mean = float(x.mean())
    std = float(x.std(ddof=1)) if len(x) > 1 else 0.0
    cv = std / mean if mean != 0 else 0.0
    return {
        "mean": mean,
        "std": std,
        "cv": cv,
        "p25": float(np.percentile(x, 25)),
        "p50": float(np.percentile(x, 50)),
        "p75": float(np.percentile(x, 75)),
        "p95": float(np.percentile(x, 95)),
    }
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_eval_metrics.py -v
```

Expected: 6 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/eval/ tests/test_eval_metrics.py
git commit -m "feat(eval): kinematics metrics (speed, accel, jerk, stats)"
```

---

## Task 13: Eval — FFT Spectrum + KL Divergence

**Files:**
- Modify: `ai_mouse/eval/metrics.py` (add FFT and KL functions)
- Modify: `tests/test_eval_metrics.py`

- [ ] **Step 1: Write failing tests**

Append to `tests/test_eval_metrics.py`:

```python
class TestFftSpectrum:
    def test_finds_dominant_frequency(self):
        """A pure 8 Hz signal should have its peak near 8 Hz."""
        from ai_mouse.eval.metrics import fft_spectrum
        # Sample at 100 Hz for 1 second
        sample_rate_hz = 100.0
        ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz)
        signal = np.sin(2 * np.pi * 8 * ts_ms / 1000)  # 8 Hz sine
        freqs, mags = fft_spectrum(signal, sample_rate_hz)
        peak_freq = freqs[np.argmax(mags)]
        assert abs(peak_freq - 8.0) < 1.0  # within 1 Hz

    def test_returns_only_positive_frequencies(self):
        from ai_mouse.eval.metrics import fft_spectrum
        signal = np.random.randn(64)
        freqs, mags = fft_spectrum(signal, 50.0)
        assert (freqs >= 0).all()
        assert len(freqs) == len(mags)


class TestKlDivergence:
    def test_identical_distributions_zero_kl(self):
        """KL(p, p) ≈ 0."""
        from ai_mouse.eval.metrics import kl_divergence_histograms
        rng = np.random.default_rng(42)
        x = rng.normal(0, 1, 5000)
        y = rng.normal(0, 1, 5000)
        kl = kl_divergence_histograms(x, y, bins=50)
        assert kl < 0.05

    def test_different_distributions_positive_kl(self):
        """Different means → positive KL."""
        from ai_mouse.eval.metrics import kl_divergence_histograms
        rng = np.random.default_rng(42)
        x = rng.normal(0, 1, 5000)
        y = rng.normal(3, 1, 5000)
        kl = kl_divergence_histograms(x, y, bins=50)
        assert kl > 0.5

    def test_handles_disjoint_supports(self):
        """No NaN even when histograms have non-overlapping bins."""
        from ai_mouse.eval.metrics import kl_divergence_histograms
        x = np.array([1.0, 1.1, 1.2, 1.3, 1.4])
        y = np.array([10.0, 10.1, 10.2, 10.3, 10.4])
        kl = kl_divergence_histograms(x, y, bins=10)
        assert np.isfinite(kl)
```

- [ ] **Step 2: Run tests, verify failure**

```bash
uv run pytest tests/test_eval_metrics.py::TestFftSpectrum tests/test_eval_metrics.py::TestKlDivergence -v
```

Expected: FAIL with `ImportError`.

- [ ] **Step 3: Implement fft_spectrum and kl_divergence_histograms**

Append to `ai_mouse/eval/metrics.py`:

```python
def fft_spectrum(
    signal: np.ndarray, sample_rate_hz: float
) -> tuple[np.ndarray, np.ndarray]:
    """Compute one-sided FFT magnitude spectrum.

    Args:
        signal: 1-D real-valued signal.
        sample_rate_hz: Sampling rate in Hz.

    Returns:
        (freqs, magnitudes) — positive frequencies only.
        Magnitudes are absolute values of complex FFT coefficients.
    """
    n = len(signal)
    if n == 0:
        return np.array([]), np.array([])
    # Zero-mean to remove DC component which dominates the spectrum
    s = signal - signal.mean()
    fft = np.fft.rfft(s)
    freqs = np.fft.rfftfreq(n, d=1.0 / sample_rate_hz)
    return freqs, np.abs(fft)


def kl_divergence_histograms(
    x: np.ndarray,
    y: np.ndarray,
    bins: int = 50,
    eps: float = 1e-10,
) -> float:
    """KL divergence KL(P_x || P_y) estimated via shared-bin histograms.

    Both arrays are histogrammed over their joint range. Empty bins get
    `eps` mass to avoid log(0) — keeps result finite even for disjoint
    supports.

    Args:
        x: samples from distribution P.
        y: samples from distribution Q (the "reference").
        bins: number of histogram bins.
        eps: smoothing constant for empty bins.

    Returns:
        scalar KL divergence (nats). Always finite, ≥ 0.
    """
    if len(x) == 0 or len(y) == 0:
        return 0.0
    lo = float(min(x.min(), y.min()))
    hi = float(max(x.max(), y.max()))
    if hi <= lo:
        return 0.0
    edges = np.linspace(lo, hi, bins + 1)
    px, _ = np.histogram(x, bins=edges, density=False)
    qy, _ = np.histogram(y, bins=edges, density=False)
    px = px.astype(float) + eps
    qy = qy.astype(float) + eps
    px /= px.sum()
    qy /= qy.sum()
    return float(np.sum(px * np.log(px / qy)))
```

- [ ] **Step 4: Run tests, verify pass**

```bash
uv run pytest tests/test_eval_metrics.py::TestFftSpectrum tests/test_eval_metrics.py::TestKlDivergence -v
```

Expected: 5 tests pass.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/eval/metrics.py tests/test_eval_metrics.py
git commit -m "feat(eval): add FFT spectrum and KL divergence metrics"
```

---

## Task 14: Eval — Report Generation

**Files:**
- Create: `ai_mouse/eval/report.py`
- Modify: `tests/test_eval_metrics.py` (add light-weight smoke test)

- [ ] **Step 1: Write a smoke test that the report module is importable and produces output**

Append to `tests/test_eval_metrics.py`:

```python
class TestReportGeneration:
    def test_generates_report_md(self, tmp_path):
        """Smoke test: build_report writes an MD file with all expected sections."""
        from ai_mouse.eval.report import build_report

        # Synthetic generated traces (3 traces, 50 points each)
        rng = np.random.default_rng(0)
        gen_traces = []
        for _ in range(3):
            xs = np.cumsum(rng.uniform(0, 5, 50))
            ys = np.cumsum(rng.uniform(-1, 1, 50))
            ts = np.cumsum(rng.uniform(5, 20, 50))
            gen_traces.append({"xs": xs, "ys": ys, "ts": ts})

        # Synthetic reference
        ref_traces = []
        for _ in range(5):
            xs = np.cumsum(rng.uniform(0, 5, 50))
            ys = np.cumsum(rng.uniform(-1, 1, 50))
            ts = np.cumsum(rng.uniform(5, 20, 50))
            ref_traces.append({"xs": xs, "ys": ys, "ts": ts})

        out_md = tmp_path / "report.md"
        build_report(
            generated_traces=gen_traces,
            reference_traces=ref_traces,
            output_md=out_md,
            tag="smoke-test",
            model_dir="/fake/model/dir",
        )
        assert out_md.exists()
        content = out_md.read_text(encoding="utf-8")
        assert "# Eval Report" in content
        assert "smoke-test" in content
        assert "速度" in content or "speed" in content.lower()
        assert "FFT" in content.upper()
        # PNG plots should exist next to MD
        plot_dir = tmp_path / "plots"
        assert plot_dir.exists()
        assert any(plot_dir.iterdir())
```

- [ ] **Step 2: Run test, verify failure**

```bash
uv run pytest tests/test_eval_metrics.py::TestReportGeneration -v
```

Expected: FAIL with `ImportError: cannot import build_report`.

- [ ] **Step 3: Implement report.py**

Create `ai_mouse/eval/report.py`:

```python
"""Markdown report generation for eval results.

Outputs a self-contained .md file with embedded PNG plots in a sibling
'plots/' directory.
"""
from __future__ import annotations

import logging
from datetime import datetime
from pathlib import Path

import numpy as np

import matplotlib
matplotlib.use("Agg")  # headless
import matplotlib.pyplot as plt  # noqa: E402

from ai_mouse.eval.metrics import (
    compute_acceleration,
    compute_jerk,
    compute_speed,
    compute_stats,
    fft_spectrum,
    kl_divergence_histograms,
)

logger = logging.getLogger(__name__)


def _aggregate_kinematics(traces: list[dict]) -> dict[str, np.ndarray]:
    """Concatenate per-trace speed/accel/jerk arrays from a list of traces.

    Args:
        traces: list of {"xs", "ys", "ts"} dicts (1-D ndarrays).

    Returns:
        dict with keys "speed", "accel", "jerk", "dt" — each a flat ndarray.
    """
    speeds, accels, jerks, dts = [], [], [], []
    for tr in traces:
        xs, ys, ts = tr["xs"], tr["ys"], tr["ts"]
        if len(xs) < 4:
            continue
        v = compute_speed(xs, ys, ts)
        a = compute_acceleration(v, ts)
        j = compute_jerk(a, ts)
        speeds.append(v)
        accels.append(a)
        jerks.append(j)
        dts.append(np.diff(ts))
    return {
        "speed": np.concatenate(speeds) if speeds else np.array([]),
        "accel": np.concatenate(accels) if accels else np.array([]),
        "jerk": np.concatenate(jerks) if jerks else np.array([]),
        "dt": np.concatenate(dts) if dts else np.array([]),
    }


def _plot_distribution(
    gen: np.ndarray,
    ref: np.ndarray,
    title: str,
    output: Path,
    xlabel: str,
    bins: int = 50,
) -> None:
    """Side-by-side histogram of gen vs ref."""
    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
    if len(gen) > 0:
        ax.hist(gen, bins=bins, alpha=0.5, label="生成", density=True)
    if len(ref) > 0:
        ax.hist(ref, bins=bins, alpha=0.5, label="参考 (Balabit)", density=True)
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel("密度")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def _plot_fft_overlay(
    gen_traces: list[dict],
    ref_traces: list[dict],
    output: Path,
    sample_rate_hz: float = 100.0,
) -> None:
    """Average FFT magnitude over lateral component for gen vs ref."""
    def _avg_spectrum(traces: list[dict]) -> tuple[np.ndarray, np.ndarray]:
        all_freqs = None
        all_mags = []
        for tr in traces:
            xs, ys = tr["xs"], tr["ys"]
            if len(xs) < 8:
                continue
            # Use the cross-track ('lateral') signal: project onto perpendicular
            # of start→end vector. Approximate by detrended y.
            sig = ys - np.linspace(ys[0], ys[-1], len(ys))
            f, m = fft_spectrum(sig, sample_rate_hz)
            if all_freqs is None:
                all_freqs = f
                all_mags.append(m)
            elif len(m) == len(all_freqs):
                all_mags.append(m)
        if not all_mags:
            return np.array([]), np.array([])
        return all_freqs, np.mean(all_mags, axis=0)

    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
    f_gen, m_gen = _avg_spectrum(gen_traces)
    f_ref, m_ref = _avg_spectrum(ref_traces)
    if len(f_gen) > 0:
        ax.plot(f_gen, m_gen, label="生成", alpha=0.7)
    if len(f_ref) > 0:
        ax.plot(f_ref, m_ref, label="参考 (Balabit)", alpha=0.7)
    ax.axvspan(4, 12, alpha=0.1, color="green", label="生理震颤区间 4–12 Hz")
    ax.set_title("FFT 频谱（横向偏移信号）")
    ax.set_xlabel("Hz")
    ax.set_ylabel("|FFT|")
    ax.set_xlim(0, sample_rate_hz / 2)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def _plot_paths_overlay(traces: list[dict], output: Path, max_traces: int = 5) -> None:
    """Plot up to N generated trajectories on the same axes."""
    fig, ax = plt.subplots(figsize=(6, 5), dpi=100)
    for i, tr in enumerate(traces[:max_traces]):
        ax.plot(tr["xs"], tr["ys"], alpha=0.6, label=f"路径 {i+1}")
    ax.invert_yaxis()  # screen coords
    ax.set_title(f"前 {min(max_traces, len(traces))} 条生成轨迹")
    ax.set_xlabel("x (px)")
    ax.set_ylabel("y (px)")
    ax.set_aspect("equal", adjustable="datalim")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def build_report(
    generated_traces: list[dict],
    reference_traces: list[dict],
    output_md: Path,
    tag: str,
    model_dir: str,
    sample_rate_hz: float = 100.0,
) -> None:
    """Build a Markdown eval report with embedded plots.

    Args:
        generated_traces: list of {"xs","ys","ts"} from the generator under test.
        reference_traces: list of {"xs","ys","ts"} from Balabit (ground truth).
        output_md: destination .md path. plots/ created in same dir.
        tag: short identifier (e.g. "baseline", "post-finetune").
        model_dir: model directory path string (for provenance).
        sample_rate_hz: nominal sample rate for FFT (mouse data is irregular —
                        100 Hz is a sensible nominal).
    """
    plot_dir = output_md.parent / "plots"
    plot_dir.mkdir(parents=True, exist_ok=True)

    gen_kin = _aggregate_kinematics(generated_traces)
    ref_kin = _aggregate_kinematics(reference_traces)

    # --- KL divergences ---
    kl_speed = kl_divergence_histograms(gen_kin["speed"], ref_kin["speed"])
    kl_accel = kl_divergence_histograms(gen_kin["accel"], ref_kin["accel"])
    kl_jerk = kl_divergence_histograms(gen_kin["jerk"], ref_kin["jerk"])
    kl_dt = kl_divergence_histograms(gen_kin["dt"], ref_kin["dt"])

    # --- Stats ---
    stats_gen = {k: compute_stats(v) for k, v in gen_kin.items()}
    stats_ref = {k: compute_stats(v) for k, v in ref_kin.items()}

    # --- Plots ---
    _plot_distribution(gen_kin["speed"], ref_kin["speed"],
                       "速度分布", plot_dir / f"{tag}-speed.png", "px/ms")
    _plot_distribution(gen_kin["accel"], ref_kin["accel"],
                       "加速度分布", plot_dir / f"{tag}-accel.png", "px/ms²")
    _plot_distribution(gen_kin["jerk"], ref_kin["jerk"],
                       "Jerk 分布", plot_dir / f"{tag}-jerk.png", "px/ms³")
    _plot_distribution(gen_kin["dt"], ref_kin["dt"],
                       "Δt 分布", plot_dir / f"{tag}-dt.png", "ms")
    _plot_fft_overlay(generated_traces, reference_traces,
                      plot_dir / f"{tag}-fft.png", sample_rate_hz)
    _plot_paths_overlay(generated_traces, plot_dir / f"{tag}-paths.png")

    # --- Markdown ---
    now = datetime.now().strftime("%Y-%m-%d %H:%M")
    lines = [
        f"# Eval Report: {tag} ({now})",
        "",
        "## 模型信息",
        f"- Checkpoint dir: `{model_dir}`",
        f"- 生成样本数: {len(generated_traces)}",
        f"- 参考样本数: {len(reference_traces)}",
        "",
        "## KL 散度（生成 vs 参考，越小越好）",
        "| 指标 | KL |",
        "|---|---|",
        f"| 速度分布 | {kl_speed:.4f} |",
        f"| 加速度分布 | {kl_accel:.4f} |",
        f"| Jerk 分布 | {kl_jerk:.4f} |",
        f"| Δt 分布 | {kl_dt:.4f} |",
        "",
        "## 摘要统计",
        "| 指标 | 生成 mean | 参考 mean | 生成 CV | 参考 CV |",
        "|---|---|---|---|---|",
    ]
    for key, label in [("speed", "速度"), ("accel", "加速度"), ("jerk", "jerk"), ("dt", "Δt")]:
        lines.append(
            f"| {label} | {stats_gen[key]['mean']:.4g} | {stats_ref[key]['mean']:.4g} | "
            f"{stats_gen[key]['cv']:.3f} | {stats_ref[key]['cv']:.3f} |"
        )
    lines += [
        "",
        "## 直方图",
        f"![速度](plots/{tag}-speed.png)",
        f"![加速度](plots/{tag}-accel.png)",
        f"![Jerk](plots/{tag}-jerk.png)",
        f"![Δt](plots/{tag}-dt.png)",
        "",
        "## FFT 频谱（横向偏移）",
        f"![FFT](plots/{tag}-fft.png)",
        "",
        "## 生成轨迹示例",
        f"![轨迹](plots/{tag}-paths.png)",
        "",
    ]
    output_md.write_text("\n".join(lines), encoding="utf-8")
    logger.info("Report written to %s", output_md)
```

- [ ] **Step 4: Run smoke test, verify pass**

```bash
uv run pytest tests/test_eval_metrics.py::TestReportGeneration -v
```

Expected: 1 test passes. Plots and Markdown file created in `tmp_path`.

- [ ] **Step 5: Commit**

```bash
git add ai_mouse/eval/report.py tests/test_eval_metrics.py
git commit -m "feat(eval): Markdown report builder with matplotlib plots"
```

---

## Task 15: Eval — CLI

**Files:**
- Create: `ai_mouse/eval/__main__.py`

- [ ] **Step 1: Implement the eval CLI**

Create `ai_mouse/eval/__main__.py`:

```python
"""CLI: `python -m ai_mouse.eval --model-dir ... --reference ... --output ...`

Loads N synthetic start/end pairs, calls the generator, loads M reference
traces from a Balabit-format jsonl, and writes a Markdown report.
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import random
import sys
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
    """Load up to n_samples reference traces from a JSONL file.

    Returns list of {"xs","ys","ts"} 1-D ndarrays.
    """
    out: list[dict] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            moves = [e for e in rec.get("events", []) if e.get("type") == "move"]
            if len(moves) < 4:
                continue
            xs = np.array([e["x"] for e in moves], dtype=float)
            ys = np.array([e["y"] for e in moves], dtype=float)
            ts = np.array([e["t"] for e in moves], dtype=float)
            out.append({"xs": xs, "ys": ys, "ts": ts})
            if len(out) >= n_samples:
                break
    return out


def _generate_n_samples(
    model_dir: str, n_samples: int, seed: int = 0
) -> list[dict]:
    """Call the project's generator N times with random start/end pairs."""
    from ai_mouse.generator import generate

    rng = random.Random(seed)
    out: list[dict] = []
    for i in range(n_samples):
        # Random start/end on a 800x600 canvas, distance 100..600 px
        sx = rng.randint(50, 750)
        sy = rng.randint(50, 550)
        angle = rng.uniform(0, 2 * math.pi)
        dist = rng.randint(100, 600)
        ex = int(sx + dist * math.cos(angle))
        ey = int(sy + dist * math.sin(angle))
        ex = max(0, min(800, ex))
        ey = max(0, min(600, ey))
        try:
            pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
        except Exception as exc:  # noqa: BLE001
            logger.warning("generate() failed at i=%d: %s", i, exc)
            continue
        # Drop click events (last 2)
        moves = pts[:-2]
        if len(moves) < 4:
            continue
        xs = np.array([p[0] for p in moves], dtype=float)
        ys = np.array([p[1] for p in moves], dtype=float)
        ts = np.array([p[2] for p in moves], dtype=float)
        out.append({"xs": xs, "ys": ys, "ts": ts})
    return out


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate eval report comparing model output to reference traces.")
    parser.add_argument("--model-dir", required=True, help="Path to trained model dir (with flow_model.pt)")
    parser.add_argument("--reference", type=Path, required=True, help="JSONL reference traces (Balabit holdout)")
    parser.add_argument("--n-samples", type=int, default=200, help="Number of generated samples")
    parser.add_argument("--n-reference", type=int, default=1000, help="Number of reference samples to load")
    parser.add_argument("--output", type=Path, required=True, help="Output Markdown file")
    parser.add_argument("--tag", default="eval", help="Tag string used in plot filenames")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    if not Path(args.model_dir).exists():
        logger.error("Model dir not found: %s", args.model_dir)
        return 2
    if not args.reference.exists():
        logger.error("Reference jsonl not found: %s", args.reference)
        return 2

    logger.info("Loading reference from %s ...", args.reference)
    ref_traces = _load_reference_jsonl(args.reference, args.n_reference)
    logger.info("Loaded %d reference traces", len(ref_traces))

    logger.info("Generating %d samples from %s ...", args.n_samples, args.model_dir)
    gen_traces = _generate_n_samples(args.model_dir, args.n_samples, seed=args.seed)
    logger.info("Generated %d valid traces", len(gen_traces))

    if not gen_traces or not ref_traces:
        logger.error("Empty trace sets — aborting")
        return 1

    from ai_mouse.eval.report import build_report
    args.output.parent.mkdir(parents=True, exist_ok=True)
    build_report(
        generated_traces=gen_traces,
        reference_traces=ref_traces,
        output_md=args.output,
        tag=args.tag,
        model_dir=args.model_dir,
    )
    logger.info("Done. Report at %s", args.output)
    return 0


if __name__ == "__main__":
    sys.exit(main())
```

- [ ] **Step 2: Smoke-test the CLI help text**

```bash
uv run python -m ai_mouse.eval --help
```

Expected: argparse help text printed.

- [ ] **Step 3: Run all tests**

```bash
uv run pytest -x
```

Expected: all tests pass.

- [ ] **Step 4: Commit**

```bash
git add ai_mouse/eval/__main__.py
git commit -m "feat(eval): CLI for generating evaluation reports"
```

---

## Task 16: Unified Train CLI

**Files:**
- Create: `ai_mouse/__main__.py`

- [ ] **Step 1: Create unified CLI dispatcher**

Create `ai_mouse/__main__.py`:

```python
"""Unified CLI: `python -m ai_mouse {train,eval,balabit-adapter}`

Subcommands dispatch to the underlying modules. This is the recommended
top-level entry; you can also call `python -m ai_mouse.eval` etc. directly.
"""
from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path


def _train_main(args: argparse.Namespace) -> int:
    from ai_mouse.trainer import train

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    train(
        data_path=Path(args.data),
        output_dir=Path(args.output),
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        seq_len=args.seq_len,
        resume_from=Path(args.resume_from) if args.resume_from else None,
    )
    return 0


def _eval_main(args: argparse.Namespace) -> int:
    from ai_mouse.eval.__main__ import main as eval_main
    # Reconstruct argv for the sub-CLI
    argv = [
        "--model-dir", args.model_dir,
        "--reference", str(args.reference),
        "--n-samples", str(args.n_samples),
        "--n-reference", str(args.n_reference),
        "--output", str(args.output),
        "--tag", args.tag,
        "--seed", str(args.seed),
    ]
    return eval_main(argv)


def _balabit_main(args: argparse.Namespace) -> int:
    from ai_mouse.data_adapters.balabit import main as bal_main
    argv = [
        "--input", str(args.input),
        "--output", str(args.output),
        "--window-ms", str(args.window_ms),
        "--min-dist", str(args.min_dist),
        "--min-events", str(args.min_events),
        "--max-span-ms", str(args.max_span_ms),
        "--max-gap-ms", str(args.max_gap_ms),
    ]
    if args.overwrite:
        argv.append("--overwrite")
    return bal_main(argv)


def main() -> int:
    p = argparse.ArgumentParser(prog="ai_mouse", description="AI Mouse trajectory toolkit")
    sub = p.add_subparsers(dest="cmd", required=True)

    # train
    pt = sub.add_parser("train", help="Train (or fine-tune) the Flow Matching model")
    pt.add_argument("--data", required=True, help="Path to traces.jsonl")
    pt.add_argument("--output", required=True, help="Output checkpoint dir")
    pt.add_argument("--epochs", type=int, default=200)
    pt.add_argument("--batch-size", type=int, default=64)
    pt.add_argument("--lr", type=float, default=3e-4)
    pt.add_argument("--seq-len", type=int, default=64)
    pt.add_argument("--resume-from", default=None, help="Checkpoint dir to resume from (for fine-tune)")
    pt.set_defaults(func=_train_main)

    # eval
    pe = sub.add_parser("eval", help="Generate evaluation report")
    pe.add_argument("--model-dir", required=True)
    pe.add_argument("--reference", type=Path, required=True)
    pe.add_argument("--n-samples", type=int, default=200)
    pe.add_argument("--n-reference", type=int, default=1000)
    pe.add_argument("--output", type=Path, required=True)
    pe.add_argument("--tag", default="eval")
    pe.add_argument("--seed", type=int, default=0)
    pe.set_defaults(func=_eval_main)

    # balabit-adapter
    pb = sub.add_parser("balabit-adapter", help="Convert Balabit dataset to traces.jsonl")
    pb.add_argument("--input", type=Path, required=True)
    pb.add_argument("--output", type=Path, default=Path("data/pretrain_traces.jsonl"))
    pb.add_argument("--window-ms", type=int, default=1200)
    pb.add_argument("--min-dist", type=int, default=50)
    pb.add_argument("--min-events", type=int, default=5)
    pb.add_argument("--max-span-ms", type=int, default=5000)
    pb.add_argument("--max-gap-ms", type=int, default=200)
    pb.add_argument("--overwrite", action="store_true")
    pb.set_defaults(func=_balabit_main)

    args = p.parse_args()
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())
```

- [ ] **Step 2: Smoke-test all three subcommands**

```bash
uv run python -m ai_mouse train --help
uv run python -m ai_mouse eval --help
uv run python -m ai_mouse balabit-adapter --help
```

Expected: each prints argparse help.

- [ ] **Step 3: Commit**

```bash
git add ai_mouse/__main__.py
git commit -m "feat: unified CLI (python -m ai_mouse {train,eval,balabit-adapter})"
```

---

## Task 17: Run Pre-Refactor Baseline Eval (Optional)

This task captures the "before" snapshot so we can measure improvement quantitatively. **Requires**: Balabit data already converted (Task 18 needs to run first if you don't have Balabit yet, OR you can use the existing 605 `traces.jsonl` as both reference and "real data" — less informative but doable).

- [ ] **Step 1: Capture baseline using existing 605 traces as reference**

If Balabit isn't downloaded yet, use the existing 605 traces as a reference (split-off subset):

```bash
# Take the existing 605 traces as reference (use what you have)
uv run python -m ai_mouse eval \
  --model-dir data/models_v2 \
  --reference data/traces.jsonl \
  --n-samples 100 \
  --n-reference 200 \
  --output data/eval_reports/2026-05-10-baseline-pre-refactor.md \
  --tag baseline-pre-refactor
```

Expected: Markdown report at `data/eval_reports/2026-05-10-baseline-pre-refactor.md` plus PNG plots. **Note the KL values — these are the "before" numbers to beat.**

- [ ] **Step 2: Inspect the report**

Open the file in a Markdown viewer or your editor and look at:
- Speed/accel/jerk KL divergence
- Δt CV (should be very low for the current template-y model)
- The path PNG (should show the high-frequency lateral jitter)

- [ ] **Step 3: Save baseline numbers**

Manually note (or copy into a comment file) the KL values for later comparison:

```bash
echo "Pre-refactor baseline (2026-05-10):" > data/eval_reports/_BASELINE_NOTES.txt
echo "  See: 2026-05-10-baseline-pre-refactor.md" >> data/eval_reports/_BASELINE_NOTES.txt
echo "  Acceptance: post-refactor KL must be < 50% of these numbers" >> data/eval_reports/_BASELINE_NOTES.txt
```

- [ ] **Step 4: Commit baseline report**

```bash
git add data/eval_reports/2026-05-10-baseline-pre-refactor.md \
        data/eval_reports/_BASELINE_NOTES.txt \
        data/eval_reports/plots/
git commit -m "docs(eval): pre-refactor baseline report"
```

---

## Task 18: Run Balabit Adapter + Pretraining

**Prerequisites:** Download Balabit dataset to `data/balabit_raw/`:

```bash
# In a separate terminal — user manual step
cd /d/code/python/side/ai_mouse/data
git clone --depth 1 https://github.com/balabit/Mouse-Dynamics-Challenge.git balabit_raw
# OR: download zip from the GitHub page and extract to data/balabit_raw/
```

Verify the structure:

```bash
ls data/balabit_raw/
# Expect to see directories like 'training_files/' or 'test_files/' containing user folders
```

- [ ] **Step 1: Convert Balabit → pretrain_traces.jsonl**

```bash
uv run python -m ai_mouse balabit-adapter \
  --input data/balabit_raw/training_files \
  --output data/pretrain_traces.jsonl \
  --overwrite
```

Expected: log shows "Wrote N segments to data/pretrain_traces.jsonl" with N typically 5,000–50,000 depending on what's in the dataset.

If N < 1000, **stop and investigate**:
- Try `--window-ms 2000` (longer window)
- Try `--min-dist 30` (shorter minimum distance)
- Inspect `data/balabit_raw/` structure — the path may be wrong (try `data/balabit_raw/test_files` instead)

- [ ] **Step 2: Spot-check converted data**

```bash
uv run python -c "
import json
with open('data/pretrain_traces.jsonl') as f:
    lines = f.readlines()
print(f'Total: {len(lines)} traces')
sample = json.loads(lines[0])
print('Sample meta:', sample['meta'])
print('Sample event count:', len(sample['events']))
print('First 3 events:', sample['events'][:3])
"
```

Expected: meta has `start`, `end`, `dist`, `angle`, `source: balabit`. Events list looks like real `move` records with sane (x, y, t).

- [ ] **Step 3: Reserve a 5% holdout for eval**

```bash
uv run python -c "
import random
from pathlib import Path
random.seed(42)
src = Path('data/pretrain_traces.jsonl')
lines = src.read_text(encoding='utf-8').splitlines()
random.shuffle(lines)
split = int(len(lines) * 0.95)
Path('data/pretrain_train.jsonl').write_text('\n'.join(lines[:split]) + '\n', encoding='utf-8')
Path('data/balabit_holdout.jsonl').write_text('\n'.join(lines[split:]) + '\n', encoding='utf-8')
print(f'Train: {split}  Holdout: {len(lines)-split}')
"
```

- [ ] **Step 4: Run pretraining**

This will take **2 hours – 2 days** depending on hardware. Run in background or overnight.

```bash
uv run python -m ai_mouse train \
  --data data/pretrain_train.jsonl \
  --output data/models_v2_pretrained \
  --epochs 200 \
  --batch-size 128 \
  --lr 3e-4
```

Expected outputs in `data/models_v2_pretrained/`:
- `flow_model.pt`
- `click_dist.json` (will have default values since Balabit has no click events — this is expected)
- `duration_dist.json`
- `train_config.json`

- [ ] **Step 5: Verify checkpoint loads cleanly**

```bash
uv run python -c "
import torch
from ai_mouse.models import TrajectoryFlowModel
m = TrajectoryFlowModel(seq_len=64)
state = torch.load('data/models_v2_pretrained/flow_model.pt', weights_only=True)
m.load_state_dict(state)
print('Pretrain checkpoint loads OK')
"
```

Expected: `Pretrain checkpoint loads OK`.

- [ ] **Step 6: Commit (do not commit the model weights — they're in .gitignore)**

```bash
# Just commit the holdout split logic via a small note
git add docs/superpowers/plans/2026-05-10-balabit-pretrain-refactor.md
git commit -m "chore: ran Balabit conversion + pretraining (artifacts in data/, not committed)"
```

---

## Task 19: Fine-tune on 605 Traces + Final Eval

- [ ] **Step 1: Run fine-tune**

```bash
uv run python -m ai_mouse train \
  --data data/traces.jsonl \
  --output data/models_v2 \
  --epochs 50 \
  --batch-size 64 \
  --lr 1e-5 \
  --resume-from data/models_v2_pretrained
```

Expected: training completes in 5–30 minutes, `data/models_v2/flow_model.pt` updated.

- [ ] **Step 2: Generate post-refactor eval report**

```bash
uv run python -m ai_mouse eval \
  --model-dir data/models_v2 \
  --reference data/balabit_holdout.jsonl \
  --n-samples 200 \
  --n-reference 1000 \
  --output data/eval_reports/2026-05-10-final.md \
  --tag final
```

- [ ] **Step 3: Compare to baseline**

Open both reports side-by-side:
- `data/eval_reports/2026-05-10-baseline-pre-refactor.md`
- `data/eval_reports/2026-05-10-final.md`

Verify against the spec's acceptance criteria:

1. ✅ **主观**：Δt 曲线在 verify 页面应该不再 5 条重合
2. ✅ **主观**：lateral 轨迹无高频锯齿（看 paths.png）
3. ✅ **量化**：speed KL < 50% of pre-refactor baseline
4. ✅ **量化**：FFT 4–12 Hz 区间出现 peak（看 fft.png）
5. ✅ **回归**：所有非废弃测试通过

- [ ] **Step 4: Run full test suite as final regression check**

```bash
uv run pytest -v
```

Expected: all tests pass.

- [ ] **Step 5: Manual verification in the UI**

```bash
uv run python main.py
```

Open http://127.0.0.1:8765 in browser, navigate to "验证效果" tab, generate 5 paths from (100, 200) to (700, 400), and compare visually to the original screenshot.

Expected:
- 5 Δt curves visibly diverge (no longer overlapping)
- Lateral trajectories smooth, no zigzag
- Average duration similar to before (still in plausible range)

- [ ] **Step 6: Commit final report**

```bash
git add data/eval_reports/2026-05-10-final.md data/eval_reports/plots/
git commit -m "docs(eval): post-refactor final eval report

Acceptance criteria met:
- speed KL <baseline*0.5
- FFT shows 4-12Hz peak
- Δt diversity restored (subjective)
- Lateral jitter eliminated (subjective)
- All tests pass"
```

---

## Self-Review Checklist

After completing all tasks, verify:

- [ ] All 19 tasks committed
- [ ] All tests pass (`uv run pytest -v`)
- [ ] Baseline report exists for comparison
- [ ] Final report shows quantitative improvement
- [ ] UI manually verified
- [ ] No new files outside the spec's "Files Changed" list
- [ ] `data/` artifacts NOT committed (verified via `git status --short`)