# ai_mouse Library Refactor 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:** Refactor `ai_mouse` from a mixed application/library into a slim ONNX-Runtime-based inference SDK (`src/ai_mouse/`), with all training/server/eval code moved to repo-internal `tools/` and pre-trained weights bundled inside the wheel.

**Architecture:** `src/`-layout package distributable via git URL; runtime depends only on `numpy + onnxruntime`. `tools/` directory holds development-only code (torch, fastapi, etc.). Public API exposes `MouseModel`/`ScrollModel` classes plus cached `generate`/`generate_scroll` helper functions.

**Tech Stack:** Python 3.12+, NumPy, ONNX Runtime, hatchling (build backend), pytest, uv. Tools-only: PyTorch, FastAPI, scipy, matplotlib.

**Spec:** [`docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md`](../specs/2026-05-11-ai-mouse-library-design.md)

---

## File Structure (target end state)

```
ai_mouse/                         (repo root)
├── src/
│   └── ai_mouse/                 # wheel content
│       ├── __init__.py
│       ├── mouse.py
│       ├── scroll.py
│       ├── _coord.py
│       ├── _postprocess.py
│       ├── _assets.py
│       ├── errors.py
│       ├── py.typed
│       └── assets/
│           ├── flow_model.onnx
│           ├── scroll_decoder.onnx
│           ├── click_dist.json
│           ├── duration_dist.json
│           ├── train_config.json
│           └── scroll_config.json
├── tools/                        # dev-only, not in wheel
│   ├── __init__.py
│   ├── __main__.py
│   ├── train.py / serve.py / export_onnx.py
│   ├── trainer.py / models.py / collector.py / config.py
│   ├── server/ / eval/ / data_adapters/
│   └── scroll/{trainer,models,collector}.py
├── tests/{unit,tools}/
├── examples/quickstart.py
├── data/ / static/ / docs/        # unchanged
├── pyproject.toml / CHANGELOG.md / README.md / CLAUDE.md
```

---

## Phase Map

| Phase | Goal | Validation |
|---|---|---|
| 0 | Capture golden tests + train scroll model | golden npz files committed |
| 1 | Move dev-only code from `ai_mouse/` to `tools/` | `python -m tools train` works; old `from ai_mouse import generate` still works |
| 2 | Switch to `src/` layout + tighten pyproject | `uv build` produces clean wheel; runtime install has no torch |
| 3 | Write ONNX exporter + commit assets | `tools/export_onnx.py` produces `.onnx` files; parity check passes |
| 4 | Rewrite library in NumPy + ORT | Golden tests pass; `import ai_mouse` works without torch |
| 5 | Docs + cleanup | README, CHANGELOG, CLAUDE.md updated; examples runnable |

---

## Phase 0: Pre-flight

### Task 0.1: Train scroll model

The repo has `data/scroll_traces.jsonl` but no trained scroll model. The current trainer in `ai_mouse/scroll/trainer.py` exists and works.

**Files:**
- Read: `ai_mouse/scroll/trainer.py`
- Output: `data/scroll_models/{scroll_model.pt, scroll_config.json}`

- [ ] **Step 1: Locate the scroll training entry point**

Run: `uv run python -c "from ai_mouse.scroll.trainer import train; help(train)"`
Confirm there's a callable `train(data_path, output_dir, ...)` with default epochs around 100.

- [ ] **Step 2: Train the scroll model**

```bash
uv run python -c "
from pathlib import Path
from ai_mouse.scroll.trainer import train
train(
    data_path=Path('data/scroll_traces.jsonl'),
    output_dir=Path('data/scroll_models'),
)
"
```

Expected: runs ~100 epochs over ~3 minutes on CPU. Loss decreasing. Writes `scroll_model.pt` and `scroll_config.json` to `data/scroll_models/`.

- [ ] **Step 3: Verify outputs exist**

Run: `ls data/scroll_models/`
Expected: `scroll_model.pt`, `scroll_config.json`.

- [ ] **Step 4: Smoke-test inference**

```bash
uv run python -c "
from ai_mouse.scroll.generator import generate_scroll
events = generate_scroll(0, 1500, mode='target', model_dir='data/scroll_models')
print(f'Got {len(events)} events; sum deltaY = {sum(e[\"deltaY\"] for e in events)}')
"
```

Expected: prints something like `Got 14 events; sum deltaY = 1480` (close to 1500).

- [ ] **Step 5: Commit the model**

```bash
git add data/scroll_models/scroll_model.pt data/scroll_models/scroll_config.json
git commit -m "chore(scroll): train initial scroll model from scroll_traces.jsonl"
```

---

### Task 0.2: Build mouse golden npz

Capture deterministic output from the current torch-based `generate()` for use in regression tests later.

**Files:**
- Create: `scripts/build_golden_mouse.py` (temporary, will be deleted after Phase 4)
- Output: `tests/unit/data/golden_mouse.npz`

- [ ] **Step 1: Ensure tests/unit/data/ exists**

```bash
mkdir -p tests/unit/data
```

- [ ] **Step 2: Create the build script**

Create `scripts/build_golden_mouse.py`:

```python
"""One-shot script to capture golden mouse trajectories from the current torch
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
in Phase 4 produces equivalent output.

Output: tests/unit/data/golden_mouse.npz
"""
from __future__ import annotations

import random
from pathlib import Path

import numpy as np
import torch

from ai_mouse import generate

CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
    ((100, 200), (900, 400)),    # horizontal 800px
    ((500, 500), (500, 100)),    # vertical 400px upward
    ((200, 600), (800, 200)),    # 720px diagonal
    ((100, 100), (130, 110)),    # very short 31px
    ((50,  50),  (1500, 900)),   # very long 1700px
    ((400, 300), (500, 300)),    # short horizontal 100px
    ((300, 300), (700, 700)),    # 45° diagonal
    ((600, 400), (200, 100)),    # reverse diagonal
]
SEEDS = (0, 1, 2, 3)


def main() -> None:
    out: dict[str, np.ndarray] = {}
    for case_idx, (start, end) in enumerate(CASES):
        for seed in SEEDS:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            pts = generate(start=start, end=end)
            out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
    out_path = Path("tests/unit/data/golden_mouse.npz")
    np.savez_compressed(out_path, **out)
    print(f"Wrote {len(out)} golden traces to {out_path}")


if __name__ == "__main__":
    main()
```

- [ ] **Step 3: Run the script**

```bash
uv run python scripts/build_golden_mouse.py
```

Expected output:
```
Wrote 32 golden traces to tests/unit/data/golden_mouse.npz
```

- [ ] **Step 4: Inspect the npz**

```bash
uv run python -c "
import numpy as np
z = np.load('tests/unit/data/golden_mouse.npz')
print('keys:', list(z.keys())[:4], '...')
print('case0_seed0 shape:', z['case0_seed0'].shape)
print('case0_seed0 first 3 rows:', z['case0_seed0'][:3])
print('case0_seed0 last 2 rows (clicks):', z['case0_seed0'][-2:])
"
```

Expected: 32 keys, each shape (66, 3) — 64 moves + 2 click events. Last two rows share x,y; t increments.

- [ ] **Step 5: Commit the golden file (not the script yet)**

```bash
git add tests/unit/data/golden_mouse.npz scripts/build_golden_mouse.py
git commit -m "test: capture mouse generate() golden output (pre-migration)"
```

---

### Task 0.3: Build scroll golden npz

**Files:**
- Create: `scripts/build_golden_scroll.py` (temporary)
- Output: `tests/unit/data/golden_scroll.npz`

- [ ] **Step 1: Create the script**

Create `scripts/build_golden_scroll.py`:

```python
"""Capture golden scroll event sequences from current torch implementation."""
from __future__ import annotations

import random
from pathlib import Path

import numpy as np
import torch

from ai_mouse import generate_scroll

CASES: list[tuple[int, int, str]] = [
    (0,    1500, "target"),
    (0,    500,  "precise"),
    (0,    5000, "fast"),
    (2000, 0,    "target"),   # upward
    (0,    800,  "precise"),
    (0,    3500, "fast"),
    (1000, 1200, "precise"),  # tiny scroll
    (0,    10000, "fast"),     # very long
]
SEEDS = (0, 1, 2, 3)


def main() -> None:
    out: dict[str, np.ndarray] = {}
    for case_idx, (start_y, end_y, mode) in enumerate(CASES):
        for seed in SEEDS:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            events = generate_scroll(start_y, end_y, mode=mode)
            arr = np.array(
                [[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
                dtype=np.int64,
            )
            out[f"case{case_idx}_seed{seed}"] = arr
    out_path = Path("tests/unit/data/golden_scroll.npz")
    np.savez_compressed(out_path, **out)
    print(f"Wrote {len(out)} scroll golden traces to {out_path}")


if __name__ == "__main__":
    main()
```

- [ ] **Step 2: Run it**

```bash
uv run python scripts/build_golden_scroll.py
```

Expected: `Wrote 32 scroll golden traces to tests/unit/data/golden_scroll.npz`

- [ ] **Step 3: Commit**

```bash
git add tests/unit/data/golden_scroll.npz scripts/build_golden_scroll.py
git commit -m "test: capture scroll generate() golden output (pre-migration)"
```

---

## Phase 1: Move dev code out of the `ai_mouse/` package

After Phase 1, `ai_mouse/` package contains ONLY inference-related modules (still torch-based for now). All training/server/collector code lives under `tools/`. The library API `from ai_mouse import generate` still works because we haven't touched it yet.

### Task 1.1: Scaffold `tools/` directory

**Files:**
- Create: `tools/__init__.py`

- [ ] **Step 1: Create tools/ and an empty __init__.py**

```bash
mkdir -p tools/scroll
touch tools/__init__.py tools/scroll/__init__.py
```

- [ ] **Step 2: Verify**

```bash
ls tools/ tools/scroll/
```

Expected: `__init__.py` in both.

- [ ] **Step 3: Commit**

```bash
git add tools/__init__.py tools/scroll/__init__.py
git commit -m "chore: scaffold tools/ directory"
```

---

### Task 1.2: Move trainer + models + utils + config to tools/

Move the torch-using mouse modules together so internal imports stay consistent within one commit.

**Files:**
- Move: `ai_mouse/trainer.py` → `tools/trainer.py`
- Move: `ai_mouse/models.py` → `tools/models.py`
- Move: `ai_mouse/utils.py` → `tools/utils.py`
- Move: `ai_mouse/config.py` → `tools/config.py`
- Modify: `ai_mouse/generator.py` (update imports)

- [ ] **Step 1: git mv the files**

```bash
git mv ai_mouse/trainer.py tools/trainer.py
git mv ai_mouse/models.py tools/models.py
git mv ai_mouse/utils.py tools/utils.py
git mv ai_mouse/config.py tools/config.py
```

- [ ] **Step 2: Update imports inside moved files**

In `tools/trainer.py`, replace:
- `from ai_mouse.config import TrainConfig` → `from tools.config import TrainConfig`
- `from ai_mouse.coord import encode_trajectory` → `from ai_mouse.coord import encode_trajectory` (unchanged — coord stays in package)
- `from ai_mouse.models import TrajectoryFlowModel` → `from tools.models import TrajectoryFlowModel`
- `from ai_mouse.utils import resample_arc` → `from tools.utils import resample_arc`

In `tools/utils.py`: no imports to change (pure numpy).
In `tools/models.py`: no imports to change (torch only).
In `tools/config.py`: no imports to change.

- [ ] **Step 3: Update `ai_mouse/generator.py` to import torch model from tools**

Find the imports near the top:

```python
from ai_mouse.config import GenerateConfig
from ai_mouse.coord import decode_trajectory
from ai_mouse.models import TrajectoryFlowModel
from ai_mouse.utils import resample_arc
```

Replace with:

```python
from ai_mouse.coord import decode_trajectory
from tools.config import GenerateConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc
```

(Note: `GenerateConfig` also moved with `config.py`. We'll yank this cross-boundary import in Phase 4 when generator.py is replaced entirely.)

- [ ] **Step 4: Verify package imports**

```bash
uv run python -c "from ai_mouse import generate; print(generate.__module__)"
```

Expected: prints `ai_mouse.generator` with no ImportError.

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

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

Expected: all pass (some test files may need import updates — see next step).

- [ ] **Step 6: Update test imports if needed**

In `tests/test_trainer.py`, `tests/test_models.py`, `tests/conftest.py`, replace:
- `from ai_mouse.trainer import ...` → `from tools.trainer import ...`
- `from ai_mouse.models import ...` → `from tools.models import ...`
- `from ai_mouse.config import TrainConfig` → `from tools.config import TrainConfig`

- [ ] **Step 7: Re-run tests**

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

Expected: all pass.

- [ ] **Step 8: Commit**

```bash
git add -A
git commit -m "refactor: move trainer/models/utils/config to tools/"
```

---

### Task 1.3: Move scroll trainer / models / collector

**Files:**
- Move: `ai_mouse/scroll/trainer.py` → `tools/scroll/trainer.py`
- Move: `ai_mouse/scroll/models.py` → `tools/scroll/models.py`
- Move: `ai_mouse/scroll/collector.py` → `tools/scroll/collector.py`
- Modify: `ai_mouse/scroll/__init__.py`, `ai_mouse/scroll/generator.py`

- [ ] **Step 1: git mv**

```bash
git mv ai_mouse/scroll/trainer.py tools/scroll/trainer.py
git mv ai_mouse/scroll/models.py tools/scroll/models.py
git mv ai_mouse/scroll/collector.py tools/scroll/collector.py
```

- [ ] **Step 2: Update imports inside moved files**

In `tools/scroll/trainer.py`:
- `from ai_mouse.scroll.models import ScrollCVAE` → `from tools.scroll.models import ScrollCVAE`
- `from ai_mouse.config import ScrollTrainConfig` → `from tools.config import ScrollTrainConfig`

In `tools/scroll/collector.py`:
- `from ai_mouse.config import SCROLL_MODES, ScrollModeConfig` → `from tools.config import SCROLL_MODES, ScrollModeConfig`

- [ ] **Step 3: Update `ai_mouse/scroll/generator.py`**

Replace `from ai_mouse.scroll.models import ScrollCVAE` with `from tools.scroll.models import ScrollCVAE`.

- [ ] **Step 4: Strip stale imports from `ai_mouse/scroll/__init__.py`**

Read current content first:

```bash
cat ai_mouse/scroll/__init__.py
```

Edit it to only re-export `generate_scroll` (the only public surface that stays in the package):

```python
"""Scroll wheel event generation (inference only)."""
from ai_mouse.scroll.generator import generate_scroll

__all__ = ["generate_scroll"]
```

- [ ] **Step 5: Update test imports**

In `tests/test_scroll_trainer.py`, `tests/test_scroll_models.py`, `tests/test_scroll_collector.py`:
- `from ai_mouse.scroll.trainer import ...` → `from tools.scroll.trainer import ...`
- `from ai_mouse.scroll.models import ...` → `from tools.scroll.models import ...`
- `from ai_mouse.scroll.collector import ...` → `from tools.scroll.collector import ...`

In `tests/conftest.py`:
- `from ai_mouse.scroll.models import ScrollCVAE` → `from tools.scroll.models import ScrollCVAE`

- [ ] **Step 6: Run scroll tests**

```bash
uv run pytest tests/test_scroll_*.py -v
```

Expected: all pass.

- [ ] **Step 7: Commit**

```bash
git add -A
git commit -m "refactor(scroll): move trainer/models/collector to tools/scroll/"
```

---

### Task 1.4: Move mouse collector

**Files:**
- Move: `ai_mouse/collector.py` → `tools/collector.py`

- [ ] **Step 1: git mv + import fix**

```bash
git mv ai_mouse/collector.py tools/collector.py
```

In `tools/collector.py`, replace any `from ai_mouse.config import ...` with `from tools.config import ...`.

- [ ] **Step 2: Search for callers**

```bash
grep -rn "from ai_mouse.collector" --include="*.py"
grep -rn "from ai_mouse import collector" --include="*.py"
```

Update each hit to `from tools.collector import ...`.

- [ ] **Step 3: Run tests touching collector**

```bash
uv run pytest tests/ -k collector -v
```

Expected: pass.

- [ ] **Step 4: Commit**

```bash
git add -A
git commit -m "refactor: move collector to tools/"
```

---

### Task 1.5: Move server/

**Files:**
- Move: `ai_mouse/server/` → `tools/server/`
- Modify: `tools/server/__init__.py` (path resolution to static/)

- [ ] **Step 1: git mv**

```bash
git mv ai_mouse/server tools/server
```

- [ ] **Step 2: Fix imports inside tools/server/**

For each file in `tools/server/` (`__init__.py`, `deps.py`, `routes_collect.py`, `routes_train.py`, `routes_verify.py`, `routes_scroll.py`), replace:
- `from ai_mouse.collector import Collector` → `from tools.collector import Collector`
- `from ai_mouse.scroll.collector import ScrollCollector` → `from tools.scroll.collector import ScrollCollector`
- `from ai_mouse.scroll.trainer import train as train_scroll` → `from tools.scroll.trainer import train as train_scroll`
- `from ai_mouse.trainer import train` → `from tools.trainer import train`
- `from ai_mouse.config import ...` → `from tools.config import ...`

Keep these unchanged (they're library API):
- `from ai_mouse import generate, generate_scroll`

- [ ] **Step 3: Fix static path resolution in `tools/server/__init__.py`**

The current code reads:

```python
_HERE = Path(__file__).resolve().parent
_STATIC_DIR = _HERE.parent.parent / "static"
```

After moving, `_HERE` is `tools/server/` so `.parent.parent` becomes the repo root — already correct. Verify by:

```bash
uv run python -c "from tools.server import create_app; app = create_app(); print('app routes:', len(app.routes))"
```

Expected: no error; prints route count.

- [ ] **Step 4: Update test imports**

In `tests/test_server.py`:
- `from ai_mouse.server import create_app` → `from tools.server import create_app`
- Any `from ai_mouse.server.X import Y` → `from tools.server.X import Y`
- `import ai_mouse.server.deps as deps_module` (if present) → `import tools.server.deps as deps_module`

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

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

Expected: pass.

- [ ] **Step 6: Commit**

```bash
git add -A
git commit -m "refactor: move server/ to tools/server/"
```

---

### Task 1.6: Move eval/ and data_adapters/

**Files:**
- Move: `ai_mouse/eval/` → `tools/eval/`
- Move: `ai_mouse/data_adapters/` → `tools/data_adapters/`

- [ ] **Step 1: git mv**

```bash
git mv ai_mouse/eval tools/eval
git mv ai_mouse/data_adapters tools/data_adapters
```

- [ ] **Step 2: Fix imports in moved files**

In `tools/eval/__main__.py`:
- `from ai_mouse.eval.report import build_report` → `from tools.eval.report import build_report`

In `tools/eval/report.py`:
- `from ai_mouse.eval.metrics import ...` → `from tools.eval.metrics import ...`

In `tools/data_adapters/__main__.py`:
- `from ai_mouse.data_adapters.balabit import main` → `from tools.data_adapters.balabit import main`

In `tools/data_adapters/balabit.py`:
- `from ai_mouse.config import BalabitAdapterConfig` → `from tools.config import BalabitAdapterConfig`

- [ ] **Step 3: Update test imports**

In `tests/test_eval_metrics.py`:
- `from ai_mouse.eval.metrics import ...` → `from tools.eval.metrics import ...`

In `tests/test_balabit_adapter.py`:
- `from ai_mouse.data_adapters.balabit import ...` → `from tools.data_adapters.balabit import ...`

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

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

Expected: pass.

- [ ] **Step 5: Commit**

```bash
git add -A
git commit -m "refactor: move eval/ and data_adapters/ to tools/"
```

---

### Task 1.7: Move CLI dispatcher

**Files:**
- Move: `ai_mouse/__main__.py` → `tools/__main__.py`
- Modify: `tools/__main__.py` (update internal subcommand wiring)

- [ ] **Step 1: git mv**

```bash
git mv ai_mouse/__main__.py tools/__main__.py
```

- [ ] **Step 2: Update internal imports**

In `tools/__main__.py`:
- `from ai_mouse.trainer import train` → `from tools.trainer import train`
- `from ai_mouse.eval.__main__ import main as eval_main` → `from tools.eval.__main__ import main as eval_main`
- `from ai_mouse.data_adapters.balabit import main as bal_main` → `from tools.data_adapters.balabit import main as bal_main`

- [ ] **Step 3: Verify CLI dispatch**

```bash
uv run python -m tools --help
```

Expected: prints help showing `train`, `eval`, `balabit-adapter` subcommands.

```bash
uv run python -m tools train --help
```

Expected: prints `train`-specific args.

- [ ] **Step 4: Commit**

```bash
git add -A
git commit -m "refactor: move CLI dispatcher to tools/__main__.py"
```

---

### Task 1.8: Convert root `main.py` to `tools/serve.py`

**Files:**
- Move: `main.py` → `tools/serve.py`

- [ ] **Step 1: git mv**

```bash
git mv main.py tools/serve.py
```

- [ ] **Step 2: Fix imports in tools/serve.py**

```python
from tools.server import create_app
```

(was `from ai_mouse.server import create_app`)

- [ ] **Step 3: Verify it starts**

In one terminal: `uv run python tools/serve.py`. In another: `curl http://127.0.0.1:8765/api/status` (or similar status endpoint). Kill the server. Then:

```bash
uv run python -c "from tools.serve import app; print('app:', app)"
```

Expected: prints `app: <FastAPI ...>` without error.

- [ ] **Step 4: Commit**

```bash
git add -A
git commit -m "refactor: move web entry main.py to tools/serve.py"
```

---

### Task 1.9: Split tests into `tests/unit/` and `tests/tools/`

**Files:**
- Move test files based on dependency:
  - `tests/unit/`: `test_coord.py`, `test_generator.py` (still uses torch via current `generate()` — KEEP in unit; will be rewritten in Phase 4)
  - `tests/tools/`: `test_trainer.py`, `test_models.py`, `test_server.py`, `test_scroll_*.py`, `test_eval_metrics.py`, `test_balabit_adapter.py`

Special case: `test_generator.py` and `test_coord.py` test the library API — they belong in `tests/unit/`. They depend on torch transitively today (via the current generator.py) but in Phase 4 they will not. Move them now to `tests/unit/`; they will keep working through both phases.

- [ ] **Step 1: Create test dirs and split**

```bash
mkdir -p tests/unit tests/tools
git mv tests/test_coord.py tests/unit/test_coord.py
git mv tests/test_generator.py tests/unit/test_generator.py
git mv tests/test_trainer.py tests/tools/test_trainer.py
git mv tests/test_models.py tests/tools/test_models.py
git mv tests/test_server.py tests/tools/test_server.py
git mv tests/test_scroll_collector.py tests/tools/test_scroll_collector.py
git mv tests/test_scroll_generator.py tests/unit/test_scroll_generator.py
git mv tests/test_scroll_models.py tests/tools/test_scroll_models.py
git mv tests/test_scroll_trainer.py tests/tools/test_scroll_trainer.py
git mv tests/test_eval_metrics.py tests/tools/test_eval_metrics.py
git mv tests/test_balabit_adapter.py tests/tools/test_balabit_adapter.py
```

- [ ] **Step 2: Split conftest.py**

Current `tests/conftest.py` provides `model_dir` and `scroll_model_dir` fixtures that use torch. These are used by tests that will end up in `tests/tools/` (the torch-using ones). Move them there:

```bash
git mv tests/conftest.py tests/tools/conftest.py
```

Create empty `tests/unit/conftest.py`:

```python
"""Fixtures for library-only tests (no torch)."""
```

- [ ] **Step 3: Add __init__.py if pytest needs them**

```bash
touch tests/unit/__init__.py tests/tools/__init__.py
```

(`tests/__init__.py` already exists.)

- [ ] **Step 4: Run both directories separately**

```bash
uv run pytest tests/unit -v
uv run pytest tests/tools -v
```

Expected: both pass. Some tests in tests/unit may still touch torch indirectly via the current generator.py — that's OK, will be cleared in Phase 4.

- [ ] **Step 5: Commit**

```bash
git add -A
git commit -m "refactor(tests): split into tests/unit and tests/tools"
```

---

### Task 1.10: Verify whole Phase 1 outcome

- [ ] **Step 1: Inspect package surface**

```bash
ls ai_mouse/
```

Expected (Phase 1 end state): `__init__.py`, `coord.py`, `generator.py`, `scroll/` (with `__init__.py`, `generator.py` only).

`ai_mouse/scroll/`:

```bash
ls ai_mouse/scroll/
```

Expected: `__init__.py`, `generator.py`.

- [ ] **Step 2: Verify imports from each side still work**

```bash
uv run python -c "
from ai_mouse import generate, generate_scroll
print('Library import OK')
from tools.trainer import train
from tools.scroll.trainer import train as st
from tools.server import create_app
from tools.eval.metrics import compute_speed
print('Tools imports OK')
"
```

Expected: prints both OK lines.

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

```bash
uv run pytest tests/ -v
```

Expected: all green.

---

## Phase 2: Switch to `src/` layout + tighten pyproject

### Task 2.1: git mv `ai_mouse` → `src/ai_mouse`

**Files:**
- Move: `ai_mouse/` → `src/ai_mouse/`

- [ ] **Step 1: Move the package**

```bash
mkdir -p src
git mv ai_mouse src/ai_mouse
```

- [ ] **Step 2: Verify nothing inside needs path updates**

The package code uses absolute imports like `from ai_mouse.coord import ...`. After the move, `ai_mouse` is still importable (because src/ becomes a path entry for setuptools/hatchling). Sanity-check there are no hard-coded paths in the source:

```bash
grep -rn "ai_mouse/" src/ai_mouse/ --include="*.py"
```

Expected: only string matches inside docstrings/comments, no live `Path("ai_mouse/...")` constructions.

- [ ] **Step 3: Commit**

```bash
git add -A
git commit -m "refactor: switch to src/ layout"
```

---

### Task 2.2: Rewrite pyproject.toml (hatchling + tightened deps)

**Files:**
- Modify: `pyproject.toml`

- [ ] **Step 1: Backup current pyproject**

```bash
cp pyproject.toml pyproject.toml.bak
```

- [ ] **Step 2: Write the new pyproject.toml**

Replace the entire file with:

```toml
[project]
name = "ai-mouse"
version = "0.2.0"
description = "Human-like mouse trajectory and scroll wheel event generator (ONNX Runtime SDK)."
requires-python = ">=3.12,<3.14"
dependencies = [
    "numpy>=1.26.0",
    "onnxruntime>=1.17.0",
]

[project.urls]
Repository = "https://github.com/<owner>/ai_mouse"

[dependency-groups]
dev = [
    "torch>=2.2.0",
    "fastapi>=0.111.0",
    "uvicorn>=0.29.0",
    "scipy>=1.10.0",
    "matplotlib>=3.8.0",
    "pytest>=8.0.0",
    "pytest-asyncio>=0.23.0",
    "httpx>=0.27.0",
    "onnx>=1.15.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/ai_mouse"]

[tool.hatch.build.targets.wheel.force-include]
"src/ai_mouse/assets" = "ai_mouse/assets"

[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
```

Notes:
- `onnx` is in `[dev]` only because the export script in `tools/export_onnx.py` uses it; runtime doesn't.
- The `force-include` line is harmless before assets/ exists; it becomes load-bearing in Phase 3.

- [ ] **Step 3: Re-sync dev environment**

```bash
uv sync --group dev
```

Expected: completes without error. `uv.lock` updates.

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

```bash
uv run pytest tests/ -v
```

Expected: all pass.

- [ ] **Step 5: Delete the backup**

```bash
rm pyproject.toml.bak
```

- [ ] **Step 6: Commit**

```bash
git add pyproject.toml uv.lock
git commit -m "build: switch to hatchling + src layout; tighten runtime deps"
```

---

### Task 2.3: Smoke-test wheel build

**Files:**
- (no files modified, just verify build)

- [ ] **Step 1: Build the wheel**

```bash
uv build
```

Expected: produces `dist/ai_mouse-0.2.0-py3-none-any.whl` and `dist/ai_mouse-0.2.0.tar.gz`.

- [ ] **Step 2: Inspect wheel contents**

```bash
uv run python -c "
import zipfile
with zipfile.ZipFile('dist/ai_mouse-0.2.0-py3-none-any.whl') as z:
    for n in z.namelist():
        print(n)
"
```

Expected: shows `ai_mouse/__init__.py`, `ai_mouse/generator.py`, `ai_mouse/coord.py`, etc. No `tools/` content; no `tests/`.

- [ ] **Step 3: Try installing into a clean venv**

```bash
uv venv .venv-clean
.venv-clean/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl
.venv-clean/Scripts/python -c "import ai_mouse; print(ai_mouse.__file__)"
```

Expected: import works. Note `torch` is NOT installed in this venv, so `from ai_mouse import generate` will FAIL right now (current generator.py still imports torch). That's expected pre-Phase-4 — just confirm `import ai_mouse` itself succeeds (it doesn't trigger generator.py).

Actually `ai_mouse/__init__.py` does `from ai_mouse.generator import generate`, which transitively imports torch. So this import WILL fail. Expected outcome:

```
ModuleNotFoundError: No module named 'torch'
```

Confirms the wheel content is correct but the runtime promise isn't met yet — exactly the state we expect at Phase 2 end. Document this in commit message.

- [ ] **Step 4: Clean up**

```bash
rm -rf .venv-clean dist/
```

- [ ] **Step 5: No commit needed (verification only)**

---

## Phase 3: ONNX export

### Task 3.1: Write the mouse-model export portion of `tools/export_onnx.py`

**Files:**
- Create: `tools/export_onnx.py`

- [ ] **Step 1: Create the file with imports and helpers**

Create `tools/export_onnx.py`:

```python
"""Export trained PyTorch checkpoints to ONNX for the inference SDK.

Usage:
    uv run python tools/export_onnx.py \
        --flow-ckpt data/models_v2 \
        --scroll-ckpt data/scroll_models \
        --output src/ai_mouse/assets/

Produces:
    <output>/flow_model.onnx
    <output>/scroll_decoder.onnx
    <output>/click_dist.json
    <output>/duration_dist.json
    <output>/train_config.json
    <output>/scroll_config.json

A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails
the .onnx files are deleted to prevent shipping broken weights.
"""
from __future__ import annotations

import argparse
import json
import logging
import shutil
import sys
from pathlib import Path

import numpy as np
import torch

logger = logging.getLogger(__name__)

_ATOL = 1e-4
```

- [ ] **Step 2: Add `export_flow_model` function**

Append to `tools/export_onnx.py`:

```python
def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export TrajectoryFlowModel to ONNX.

    Args:
        ckpt_dir: directory with flow_model.pt and train_config.json.
        out_dir: destination directory (created if missing).

    Returns:
        Path to the written flow_model.onnx.
    """
    from tools.models import TrajectoryFlowModel

    config_path = ckpt_dir / "train_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    d_model = int(cfg["d_model"])
    nhead = int(cfg["nhead"])
    num_layers = int(cfg["num_layers"])
    dim_feedforward = int(cfg["dim_feedforward"])
    cond_dim = int(cfg.get("cond_dim", 3))

    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        dim_feedforward=dim_feedforward,
        cond_dim=cond_dim,
        dropout=0.0,  # disable dropout for export
    )
    state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "flow_model.onnx"

    dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32)
    dummy_t = torch.zeros(1, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        model,
        (dummy_x, dummy_t, dummy_cond),
        str(out_path),
        input_names=["x_t", "t", "cond"],
        output_names=["v"],
        dynamic_axes={
            "x_t": {0: "batch"},
            "t": {0: "batch"},
            "cond": {0: "batch"},
            "v": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
    )
    logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6)
    return out_path
```

- [ ] **Step 3: Add a quick sanity test (manual run)**

In a python shell:

```bash
uv run python -c "
from pathlib import Path
from tools.export_onnx import export_flow_model
out = export_flow_model(Path('data/models_v2'), Path('/tmp/test_export'))
print('Wrote:', out)
"
```

Expected: prints the output path and a size like 2-3 MB. The file exists.

- [ ] **Step 4: Commit**

```bash
git add tools/export_onnx.py
git commit -m "feat(tools): add export_flow_model for ONNX export"
```

---

### Task 3.2: Add scroll-decoder export

**Files:**
- Modify: `tools/export_onnx.py`

- [ ] **Step 1: Define ScrollDecoder wrapper module**

Append to `tools/export_onnx.py`:

```python
class _ScrollDecoder(torch.nn.Module):
    """Wraps ScrollCVAE.decode for ONNX export.

    The full ScrollCVAE is encoder+decoder; inference only needs decoder.
    """

    def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int):
        super().__init__()
        self.dec_h0 = dec_h0
        self.dec_gru = dec_gru
        self.dec_out = dec_out
        self.seq_len = seq_len
        self.hidden = hidden

    def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        b = z.shape[0]
        zc = torch.cat([z, cond], dim=-1)
        h0_flat = self.dec_h0(zc)
        h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous()
        inp = zc.unsqueeze(1).expand(b, self.seq_len, -1)
        out, _ = self.dec_gru(inp, h0)
        return self.dec_out(out)
```

- [ ] **Step 2: Add `export_scroll_decoder` function**

Append:

```python
def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export ScrollCVAE decoder to ONNX."""
    from tools.scroll.models import ScrollCVAE

    config_path = ckpt_dir / "scroll_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    latent_dim = int(cfg["latent_dim"])
    hidden = int(cfg["hidden"])
    cond_dim = int(cfg["cond_dim"])

    full = ScrollCVAE(
        seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim
    )
    state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
    full.load_state_dict(state)
    full.eval()

    decoder = _ScrollDecoder(
        dec_h0=full.dec_h0,
        dec_gru=full.dec_gru,
        dec_out=full.dec_out,
        seq_len=seq_len,
        hidden=hidden,
    )
    decoder.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "scroll_decoder.onnx"

    dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        decoder,
        (dummy_z, dummy_cond),
        str(out_path),
        input_names=["z", "cond"],
        output_names=["seq"],
        dynamic_axes={
            "z": {0: "batch"},
            "cond": {0: "batch"},
            "seq": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
    )
    logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3)
    return out_path
```

- [ ] **Step 3: Manual sanity test**

```bash
uv run python -c "
from pathlib import Path
from tools.export_onnx import export_scroll_decoder
out = export_scroll_decoder(Path('data/scroll_models'), Path('/tmp/test_export'))
print('Wrote:', out)
"
```

Expected: prints path; file <300 KB.

- [ ] **Step 4: Commit**

```bash
git add tools/export_onnx.py
git commit -m "feat(tools): add export_scroll_decoder for ONNX export"
```

---

### Task 3.3: Add PyTorch vs ORT parity check

**Files:**
- Modify: `tools/export_onnx.py`

- [ ] **Step 1: Add parity helpers**

Append to `tools/export_onnx.py`:

```python
def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None:
    """Verify ONNX flow model matches PyTorch output on random input."""
    import onnxruntime as ort
    from tools.models import TrajectoryFlowModel

    cfg = json.loads((ckpt_dir / "train_config.json").read_text())
    seq_len = int(cfg["seq_len"])
    cond_dim = int(cfg.get("cond_dim", 3))

    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=int(cfg["d_model"]),
        nhead=int(cfg["nhead"]),
        num_layers=int(cfg["num_layers"]),
        dim_feedforward=int(cfg["dim_feedforward"]),
        cond_dim=cond_dim,
        dropout=0.0,
    )
    model.load_state_dict(
        torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
    )
    model.eval()

    torch.manual_seed(42)
    np.random.seed(42)
    x = torch.randn(2, seq_len, 3, dtype=torch.float32)
    t = torch.tensor([0.0, 0.5], dtype=torch.float32)
    cond = torch.randn(2, cond_dim, dtype=torch.float32)

    with torch.no_grad():
        torch_out = model(x, t, cond).numpy()

    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    ort_out = sess.run(
        ["v"],
        {
            "x_t": x.numpy(),
            "t": t.numpy(),
            "cond": cond.numpy(),
        },
    )[0]

    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
        max_diff = float(np.abs(torch_out - ort_out).max())
        raise RuntimeError(
            f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
        )
    logger.info("Flow model parity OK (atol=%.0e)", _ATOL)


def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None:
    """Verify ONNX scroll decoder matches PyTorch decoder output."""
    import onnxruntime as ort
    from tools.scroll.models import ScrollCVAE

    cfg = json.loads((ckpt_dir / "scroll_config.json").read_text())
    seq_len = int(cfg["seq_len"])
    latent_dim = int(cfg["latent_dim"])
    cond_dim = int(cfg["cond_dim"])

    full = ScrollCVAE(
        seq_len=seq_len,
        latent_dim=latent_dim,
        hidden=int(cfg["hidden"]),
        cond_dim=cond_dim,
    )
    full.load_state_dict(
        torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
    )
    full.eval()

    torch.manual_seed(7)
    z = torch.randn(2, latent_dim, dtype=torch.float32)
    cond = torch.randn(2, cond_dim, dtype=torch.float32)

    with torch.no_grad():
        torch_out = full.decode(z, cond).numpy()

    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0]

    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
        max_diff = float(np.abs(torch_out - ort_out).max())
        raise RuntimeError(
            f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
        )
    logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL)
```

- [ ] **Step 2: Manual test the checks**

```bash
uv run python -c "
from pathlib import Path
from tools.export_onnx import (
    export_flow_model, export_scroll_decoder,
    _check_flow_parity, _check_scroll_parity,
)
import logging; logging.basicConfig(level=logging.INFO)
out = Path('/tmp/test_export')
export_flow_model(Path('data/models_v2'), out)
_check_flow_parity(Path('data/models_v2'), out / 'flow_model.onnx')
export_scroll_decoder(Path('data/scroll_models'), out)
_check_scroll_parity(Path('data/scroll_models'), out / 'scroll_decoder.onnx')
"
```

Expected: prints two "parity OK" lines, no exceptions.

- [ ] **Step 3: Commit**

```bash
git add tools/export_onnx.py
git commit -m "feat(tools): add ORT vs PyTorch parity check for exports"
```

---

### Task 3.4: Add CLI `main()` to `tools/export_onnx.py`

**Files:**
- Modify: `tools/export_onnx.py`

- [ ] **Step 1: Add main() and __main__ guard**

Append to `tools/export_onnx.py`:

```python
def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None:
    """Copy JSON metadata files alongside the ONNX models."""
    for name in ("click_dist.json", "duration_dist.json", "train_config.json"):
        src = flow_dir / name
        if not src.exists():
            raise FileNotFoundError(f"Required metadata missing: {src}")
        shutil.copy2(src, out_dir / name)
    src = scroll_dir / "scroll_config.json"
    if not src.exists():
        raise FileNotFoundError(f"Required metadata missing: {src}")
    shutil.copy2(src, out_dir / "scroll_config.json")


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0])
    p.add_argument("--flow-ckpt", type=Path, required=True)
    p.add_argument("--scroll-ckpt", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    args = p.parse_args(argv)

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

    args.output.mkdir(parents=True, exist_ok=True)

    flow_onnx = export_flow_model(args.flow_ckpt, args.output)
    scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output)

    try:
        _check_flow_parity(args.flow_ckpt, flow_onnx)
        _check_scroll_parity(args.scroll_ckpt, scroll_onnx)
    except RuntimeError as exc:
        logger.error("Parity check failed: %s", exc)
        flow_onnx.unlink(missing_ok=True)
        scroll_onnx.unlink(missing_ok=True)
        return 1

    _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output)
    logger.info("Export complete: %s", args.output)
    return 0


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

- [ ] **Step 2: Run the full export to produce assets**

```bash
mkdir -p src/ai_mouse/assets
uv run python tools/export_onnx.py \
    --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models \
    --output src/ai_mouse/assets/
```

Expected output (timestamps elided):
```
... INFO Wrote src/ai_mouse/assets/flow_model.onnx (2.x MB)
... INFO Wrote src/ai_mouse/assets/scroll_decoder.onnx (0.x KB)
... INFO Flow model parity OK (atol=1e-04)
... INFO Scroll decoder parity OK (atol=1e-04)
... INFO Export complete: src/ai_mouse/assets
```

- [ ] **Step 3: Verify assets directory**

```bash
ls src/ai_mouse/assets/
```

Expected: `flow_model.onnx`, `scroll_decoder.onnx`, `click_dist.json`, `duration_dist.json`, `train_config.json`, `scroll_config.json`.

- [ ] **Step 4: Commit assets + script main()**

```bash
git add tools/export_onnx.py src/ai_mouse/assets/
git commit -m "feat: export ONNX weights and metadata into src/ai_mouse/assets/"
```

---

### Task 3.5: Test ONNX export with toy model

**Files:**
- Create: `tests/tools/test_export_onnx.py`

- [ ] **Step 1: Write the failing test**

Create `tests/tools/test_export_onnx.py`:

```python
"""Validate tools.export_onnx with a tiny synthetic model."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest
import torch

from tools.export_onnx import (
    _check_flow_parity,
    _check_scroll_parity,
    export_flow_model,
    export_scroll_decoder,
)


@pytest.fixture
def tiny_flow_ckpt(tmp_path: Path) -> Path:
    """A flow model with seq_len=8, d_model=16, 1 layer — small but valid."""
    from tools.models import TrajectoryFlowModel

    cfg = {
        "seq_len": 8,
        "d_model": 16,
        "nhead": 2,
        "num_layers": 1,
        "dim_feedforward": 32,
        "cond_dim": 3,
    }
    model = TrajectoryFlowModel(**cfg, dropout=0.0)
    model.eval()
    out = tmp_path / "flow_ckpt"
    out.mkdir()
    torch.save(model.state_dict(), out / "flow_model.pt")
    (out / "train_config.json").write_text(json.dumps(cfg))
    return out


@pytest.fixture
def tiny_scroll_ckpt(tmp_path: Path) -> Path:
    """A scroll model with seq_len=4, latent=4, hidden=8."""
    from tools.scroll.models import ScrollCVAE

    cfg = {"seq_len": 4, "latent_dim": 4, "hidden": 8, "cond_dim": 7}
    model = ScrollCVAE(**cfg)
    model.eval()
    out = tmp_path / "scroll_ckpt"
    out.mkdir()
    torch.save(model.state_dict(), out / "scroll_model.pt")
    (out / "scroll_config.json").write_text(json.dumps(cfg))
    return out


def test_export_flow_model_parity(tiny_flow_ckpt: Path, tmp_path: Path) -> None:
    out_dir = tmp_path / "out"
    onnx_path = export_flow_model(tiny_flow_ckpt, out_dir)
    assert onnx_path.exists()
    _check_flow_parity(tiny_flow_ckpt, onnx_path)  # raises on failure


def test_export_scroll_decoder_parity(tiny_scroll_ckpt: Path, tmp_path: Path) -> None:
    out_dir = tmp_path / "out"
    onnx_path = export_scroll_decoder(tiny_scroll_ckpt, out_dir)
    assert onnx_path.exists()
    _check_scroll_parity(tiny_scroll_ckpt, onnx_path)
```

- [ ] **Step 2: Run the tests**

```bash
uv run pytest tests/tools/test_export_onnx.py -v
```

Expected: both pass.

- [ ] **Step 3: Commit**

```bash
git add tests/tools/test_export_onnx.py
git commit -m "test(tools): cover export_onnx with tiny synthetic models"
```

---

## Phase 4: Rewrite library in NumPy + ORT

### Task 4.1: Create `_coord.py` (private numpy coordinate transforms)

**Files:**
- Create: `src/ai_mouse/_coord.py`
- Keep (for now): `src/ai_mouse/coord.py` — tools/ still imports it; deleted at end of Phase 4

- [ ] **Step 1: Copy coord.py to _coord.py**

```bash
cp src/ai_mouse/coord.py src/ai_mouse/_coord.py
```

- [ ] **Step 2: No content edits needed**

The file is already pure numpy. Verify:

```bash
grep -E "^import|^from" src/ai_mouse/_coord.py
```

Expected: only `import math` and `import numpy as np`.

- [ ] **Step 3: Write the test**

Create `tests/unit/test__coord.py`:

```python
"""Test the private numpy coordinate transforms."""
from __future__ import annotations

import numpy as np

from ai_mouse._coord import decode_trajectory, encode_trajectory


def test_encode_decode_roundtrip() -> None:
    points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]])
    start = (100, 200)
    end = (500, 300)
    encoded = encode_trajectory(points, start, end)
    decoded = decode_trajectory(encoded, start, end)
    assert np.allclose(decoded, points, atol=1e-6)


def test_encode_endpoints() -> None:
    """Start should encode to (0,0); end should encode to (1,0)."""
    points = np.array([[100.0, 200.0], [500.0, 300.0]])
    encoded = encode_trajectory(points, (100, 200), (500, 300))
    assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6)
    assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6)


def test_zero_distance_returns_zeros() -> None:
    points = np.array([[100.0, 200.0]])
    encoded = encode_trajectory(points, (100, 200), (100, 200))
    assert encoded.shape == (1, 2)
    assert np.all(encoded == 0)
```

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

```bash
uv run pytest tests/unit/test__coord.py -v
```

Expected: 3 pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_coord.py tests/unit/test__coord.py
git commit -m "feat(lib): add private _coord.py with numpy transforms"
```

---

### Task 4.2: Create `errors.py`

**Files:**
- Create: `src/ai_mouse/errors.py`

- [ ] **Step 1: Write the failing test**

Create `tests/unit/test_errors.py`:

```python
"""Test the error hierarchy."""
from __future__ import annotations

import pytest

from ai_mouse import errors


def test_model_load_error_is_aimouse_error() -> None:
    assert issubclass(errors.ModelLoadError, errors.AiMouseError)


def test_generation_error_is_aimouse_error() -> None:
    assert issubclass(errors.GenerationError, errors.AiMouseError)


def test_can_catch_specific_with_general() -> None:
    with pytest.raises(errors.AiMouseError):
        raise errors.ModelLoadError("test")
```

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

```bash
uv run pytest tests/unit/test_errors.py -v
```

Expected: ImportError on `from ai_mouse import errors`.

- [ ] **Step 3: Create the module**

Create `src/ai_mouse/errors.py`:

```python
"""Exception hierarchy for the ai_mouse library.

Downstream consumers can catch the umbrella :class:`AiMouseError`
or the specific subclasses for finer control.
"""
from __future__ import annotations


class AiMouseError(Exception):
    """Base class for all ai_mouse errors."""


class ModelLoadError(AiMouseError):
    """Raised when ONNX weights / metadata cannot be loaded."""


class GenerationError(AiMouseError):
    """Raised when inference produces an invalid result (e.g. NaN)."""
```

- [ ] **Step 4: Run test, observe pass**

```bash
uv run pytest tests/unit/test_errors.py -v
```

Expected: 3 pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/errors.py tests/unit/test_errors.py
git commit -m "feat(lib): add errors module"
```

---

### Task 4.3: Create `_assets.py` (importlib.resources loader)

**Files:**
- Create: `src/ai_mouse/_assets.py`

- [ ] **Step 1: Write the failing test**

Create `tests/unit/test_assets.py`:

```python
"""Test the asset path resolver."""
from __future__ import annotations

import json
from pathlib import Path

import pytest

from ai_mouse import _assets
from ai_mouse.errors import ModelLoadError


def test_bundled_flow_model_exists() -> None:
    p = _assets.bundled_path("flow_model.onnx")
    assert p.exists()
    assert p.suffix == ".onnx"


def test_bundled_train_config_loadable() -> None:
    p = _assets.bundled_path("train_config.json")
    cfg = json.loads(p.read_text())
    assert "seq_len" in cfg
    assert "d_model" in cfg


def test_resolve_with_custom_dir(tmp_path: Path) -> None:
    (tmp_path / "flow_model.onnx").write_bytes(b"x")
    p = _assets.resolve(tmp_path, "flow_model.onnx")
    assert p == tmp_path / "flow_model.onnx"


def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None:
    with pytest.raises(ModelLoadError, match="missing"):
        _assets.resolve(tmp_path, "nonexistent.onnx")
```

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

```bash
uv run pytest tests/unit/test_assets.py -v
```

Expected: ImportError on `from ai_mouse import _assets`.

- [ ] **Step 3: Create the module**

Create `src/ai_mouse/_assets.py`:

```python
"""Asset path resolution for bundled ONNX weights and JSON metadata.

Uses :mod:`importlib.resources` to locate files inside the installed
package, falling back to a user-supplied directory if provided.
"""
from __future__ import annotations

from importlib.resources import as_file, files
from pathlib import Path

from ai_mouse.errors import ModelLoadError

_PACKAGE_ASSETS = "ai_mouse.assets"


def bundled_path(name: str) -> Path:
    """Return a filesystem path to a bundled asset.

    Args:
        name: filename inside the assets/ directory.

    Returns:
        A concrete :class:`pathlib.Path`. Note: for zipapp installations
        this materialises a temp file; for normal site-packages installs
        it points into the package directly.
    """
    ref = files(_PACKAGE_ASSETS) / name
    # as_file is the canonical way; for non-zip installs this is a no-op
    # context that yields the actual path.
    with as_file(ref) as p:
        # We're inside the with-block; the contextmanager keeps the
        # temp file alive only while open. For zip installs we'd need
        # to extract to a stable location. For now, all our installs
        # are wheel-based (non-zip), so the path is stable after exit.
        return Path(p)


def resolve(model_path: Path | None, filename: str) -> Path:
    """Locate an asset given an optional user-supplied directory.

    Args:
        model_path: user-supplied directory, or None to use bundled assets.
        filename: file to locate inside the directory.

    Returns:
        Absolute path to the asset.

    Raises:
        ModelLoadError: if the file does not exist.
    """
    if model_path is None:
        p = bundled_path(filename)
    else:
        p = Path(model_path) / filename
    if not p.exists():
        raise ModelLoadError(f"Required asset missing: {p}")
    return p
```

- [ ] **Step 4: Run test, observe pass**

```bash
uv run pytest tests/unit/test_assets.py -v
```

Expected: 4 pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_assets.py tests/unit/test_assets.py
git commit -m "feat(lib): add _assets module for bundled-weight resolution"
```

---

### Task 4.4: Create `_postprocess.py` skeleton + `gaussian_smooth`

**Files:**
- Create: `src/ai_mouse/_postprocess.py`
- Create: `tests/unit/test_postprocess.py`

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

Create `tests/unit/test_postprocess.py`:

```python
"""Tests for trajectory post-processing primitives."""
from __future__ import annotations

import numpy as np

from ai_mouse._postprocess import gaussian_smooth


def test_gaussian_smooth_preserves_endpoints() -> None:
    x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert result[0] == 1.0
    assert result[-1] == 4.0


def test_gaussian_smooth_short_input_unchanged() -> None:
    x = np.array([1.0, 2.0, 3.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert np.array_equal(result, x)


def test_gaussian_smooth_constant_unchanged() -> None:
    x = np.full(20, 7.5)
    result = gaussian_smooth(x, sigma=1.0)
    assert np.allclose(result, x, atol=1e-6)
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

Expected: ImportError.

- [ ] **Step 3: Create module + function**

Create `src/ai_mouse/_postprocess.py`:

```python
"""Pure-numpy post-processing primitives for trajectory generation.

All functions are pure (no I/O, no global state) and accept an explicit
:class:`numpy.random.Generator` when randomness is involved.
"""
from __future__ import annotations

import numpy as np


def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-tap gaussian smoothing along a 1-D array; endpoints preserved.

    Args:
        x: 1-D input array.
        sigma: gaussian std. 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]`` unchanged.
        If ``len(x) < 5`` returns a copy of ``x`` (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()
    padded = np.pad(x, pad_width=2, mode="edge")
    smoothed = np.convolve(padded, kernel, mode="valid")
    smoothed[0] = x[0]
    smoothed[-1] = x[-1]
    return smoothed
```

- [ ] **Step 4: Run, observe pass**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

Expected: 3 pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add gaussian_smooth to _postprocess"
```

---

### Task 4.5: Add `snap_endpoints`

**Files:**
- Modify: `src/ai_mouse/_postprocess.py`
- Modify: `tests/unit/test_postprocess.py`

- [ ] **Step 1: Write failing test (append)**

Append to `tests/unit/test_postprocess.py`:

```python
from ai_mouse._postprocess import snap_endpoints


def test_snap_endpoints_pins_first_and_last() -> None:
    forward = np.linspace(0.1, 0.9, 16)
    lateral = np.full(16, 0.5)
    f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16)
    assert f[0] == 0.0
    assert l[0] == 0.0
    assert f[-1] == 1.0
    assert l[-1] == 0.0


def test_snap_endpoints_preserves_middle() -> None:
    forward = np.linspace(0.0, 1.0, 16)
    lateral = np.zeros(16)
    f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4)
    # Points before the last n_snap should be unchanged
    assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_postprocess.py::test_snap_endpoints_pins_first_and_last -v
```

Expected: ImportError.

- [ ] **Step 3: Implement**

Append to `src/ai_mouse/_postprocess.py`:

```python
def snap_endpoints(
    forward: np.ndarray,
    lateral: np.ndarray,
    seq_len: int,
    n_snap: int = 6,
) -> tuple[np.ndarray, np.ndarray]:
    """Force first point to (0,0) and last point to (1,0) with quadratic ease.

    The last ``n_snap`` points are linearly interpolated towards (1, 0)
    with quadratic easing, then the first/last points are pinned exactly.

    Args:
        forward: (T,) forward coordinates (modified in place).
        lateral: (T,) lateral coordinates (modified in place).
        seq_len: length of forward/lateral.
        n_snap: number of trailing points to ease (capped at seq_len//4).

    Returns:
        ``(forward, lateral)`` after modification.
    """
    n_snap = min(n_snap, seq_len // 4)
    for i in range(n_snap):
        alpha = ((i + 1) / n_snap) ** 2
        k = seq_len - n_snap + i
        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
    forward[0], lateral[0] = 0.0, 0.0
    forward[-1], lateral[-1] = 1.0, 0.0
    return forward, lateral
```

- [ ] **Step 4: Run all postprocess tests**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

Expected: all pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add snap_endpoints to _postprocess"
```

---

### Task 4.6: Add `smooth_start`, `enforce_forward_monotonic`

**Files:**
- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py`

- [ ] **Step 1: Write tests (append)**

```python
from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start


def test_smooth_start_dampens_lateral() -> None:
    forward = np.linspace(0, 1, 16)
    lateral = np.full(16, 1.0)
    forward[0] = lateral[0] = 0.0  # invariant: snap already done
    _, l = smooth_start(forward.copy(), lateral.copy(), n=4)
    # Lateral at points 1-4 should be < original (dampened)
    assert l[1] < 1.0
    assert l[4] < 1.0
    # Lateral at point 5+ unchanged
    assert l[5] == 1.0


def test_enforce_forward_monotonic_repairs_inversions() -> None:
    f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
    out = enforce_forward_monotonic(f.copy())
    assert np.all(np.diff(out) > 0), out


def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
    f = np.array([-0.1, 0.5, 1.2])
    out = enforce_forward_monotonic(f.copy())
    assert out[0] == 0.0
    assert out[-1] == 1.0
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

- [ ] **Step 3: Implement (append to _postprocess.py)**

```python
def smooth_start(
    forward: np.ndarray,
    lateral: np.ndarray,
    n: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
    """Dampen lateral oscillation in the first ``n`` points.

    Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
    forced non-decreasing locally; lateral is linearly damped towards 0.
    """
    n_start_fix = min(n, len(forward) // 4)
    for i in range(1, n_start_fix + 1):
        blend = i / (n_start_fix + 1)
        forward[i] = max(forward[i], forward[i - 1])
        lateral[i] = lateral[i] * blend
    return forward, lateral


def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
    """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
    seq_len = len(forward)
    for i in range(1, seq_len - 1):
        if forward[i] < forward[i - 1]:
            forward[i] = forward[i - 1] + 0.001
    forward = np.clip(forward, 0.0, 1.0)
    forward[0] = 0.0
    forward[-1] = 1.0
    return forward
```

- [ ] **Step 4: Test**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

Expected: all pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add smooth_start, enforce_forward_monotonic"
```

---

### Task 4.7: Add `resample_arc`, `build_timestamps`

**Files:**
- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py`

- [ ] **Step 1: Tests**

Append:

```python
from ai_mouse._postprocess import build_timestamps, resample_arc


def test_resample_arc_identity_when_same_length() -> None:
    pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]])
    out = resample_arc(pts, 4)
    assert np.allclose(out, pts, atol=1e-6)


def test_resample_arc_changes_length() -> None:
    pts = np.array([[float(i), 0.0] for i in range(10)])
    out = resample_arc(pts, 5)
    assert out.shape == (5, 2)
    # Endpoints preserved
    assert np.allclose(out[0], pts[0])
    assert np.allclose(out[-1], pts[-1])


def test_build_timestamps_strictly_increasing() -> None:
    log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0])
    ts = build_timestamps(log_dt, total_duration_ms=200.0)
    assert ts[0] == 0
    assert np.all(np.diff(ts) >= 1)  # at least 1 ms apart


def test_build_timestamps_total_close_to_target() -> None:
    log_dt = np.array([1.0] * 10)
    ts = build_timestamps(log_dt, total_duration_ms=300.0)
    # Last timestamp should be roughly total - one slot
    assert abs(ts[-1] - 270) < 60  # tolerant of clipping
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_postprocess.py::test_resample_arc_identity_when_same_length -v
```

- [ ] **Step 3: Implement**

Append to `_postprocess.py`:

```python
def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
    """Resample a 2-D polyline to ``n_points`` along cumulative arc length."""
    arc = np.concatenate(
        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
    )
    s_new = np.linspace(0, arc[-1], n_points)
    return np.stack(
        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
        axis=1,
    )


def build_timestamps(
    log_dt: np.ndarray,
    total_duration_ms: float,
    dt_clip: tuple[float, float] = (2.0, 150.0),
) -> np.ndarray:
    """Convert per-step log_dt + total duration to cumulative ms timestamps.

    Args:
        log_dt: (N,) array of natural-log step intervals.
        total_duration_ms: target total span. The output is scaled so the
            sum approximately matches this (modulo dt_clip).
        dt_clip: (min, max) per-step clamp in milliseconds.

    Returns:
        (N,) integer-rounded cumulative timestamps starting at 0,
        strictly increasing.
    """
    n = len(log_dt)
    dt_raw = np.clip(np.exp(log_dt), 0.0, None)
    dt_sum = dt_raw.sum()
    if dt_sum > 1e-6:
        scale = total_duration_ms / dt_sum
    else:
        scale = total_duration_ms / max(n, 1)
    dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1])

    t_abs = np.cumsum(dt_ms)
    t_abs = np.concatenate([[0.0], t_abs[:-1]])

    for i in range(1, n):
        if t_abs[i] <= t_abs[i - 1]:
            t_abs[i] = t_abs[i - 1] + 1.0
    return t_abs
```

- [ ] **Step 4: Run**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add resample_arc, build_timestamps"
```

---

### Task 4.8: Add `sample_duration` + `truncnorm_sample`

**Files:**
- Modify: `src/ai_mouse/_postprocess.py`, `tests/unit/test_postprocess.py`

- [ ] **Step 1: Tests**

```python
from ai_mouse._postprocess import sample_duration, truncnorm_sample


def test_truncnorm_sample_within_bounds() -> None:
    rng = np.random.default_rng(0)
    samples = [
        truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500)
    ]
    arr = np.array(samples)
    assert arr.min() >= 20.0
    assert arr.max() <= 300.0
    # Mean roughly close to mu
    assert abs(arr.mean() - 80.0) < 5.0


def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None:
    rng = np.random.default_rng(0)
    # mu far outside [low, high] — rejection will fail
    v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng)
    assert 20.0 <= v <= 30.0


def test_sample_duration_uses_correct_bin() -> None:
    dist_dict = {
        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
        "params": [
            {"mu_log": 4.0, "sigma_log": 0.01},  # bin 0: dist < 50
            {"mu_log": 5.0, "sigma_log": 0.01},  # bin 1: 50 <= dist < 100
            {"mu_log": 6.0, "sigma_log": 0.01},  # bin 2: 100 <= dist < 200
        ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5,
    }
    rng = np.random.default_rng(0)
    v = sample_duration(dist_dict, 150.0, rng)
    # exp(6) ~ 403, with tiny sigma we should land near there
    assert 350 < v < 460
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

- [ ] **Step 3: Implement**

Append to `_postprocess.py`:

```python
def truncnorm_sample(
    mu: float,
    sigma: float,
    low: float,
    high: float,
    rng: np.random.Generator,
    max_tries: int = 32,
) -> float:
    """Sample from N(mu, sigma) truncated to [low, high] via rejection.

    Falls back to clipping if rejection fails ``max_tries`` times.
    """
    for _ in range(max_tries):
        v = rng.normal(mu, sigma)
        if low <= v <= high:
            return float(v)
    return float(np.clip(rng.normal(mu, sigma), low, high))


def sample_duration(
    duration_dist: dict,
    dist: float,
    rng: np.random.Generator,
) -> float:
    """Sample total trajectory duration (ms) for the given pixel distance.

    Uses per-bin log-normal parameters in ``duration_dist``.
    """
    bins = duration_dist["bins"]
    params = duration_dist["params"]
    bin_idx = len(bins) - 1
    for i in range(len(bins) - 1):
        if dist < bins[i + 1]:
            bin_idx = i
            break
    bin_idx = min(bin_idx, len(params) - 1)
    mu_log = params[bin_idx]["mu_log"]
    sigma_log = params[bin_idx]["sigma_log"]
    return float(np.exp(rng.normal(mu_log, sigma_log)))
```

- [ ] **Step 4: Test**

```bash
uv run pytest tests/unit/test_postprocess.py -v
```

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add sample_duration, truncnorm_sample (no scipy)"
```

---

### Task 4.9: Write `mouse.py` (`MouseModel` + `_get_default_mouse_model`)

**Files:**
- Create: `src/ai_mouse/mouse.py`

- [ ] **Step 1: Write test scaffolding**

Create `tests/unit/test_mouse.py`:

```python
"""Tests for MouseModel and ai_mouse.generate()."""
from __future__ import annotations

import numpy as np
import pytest


def test_mouse_model_init_default() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    assert m._seq_len > 0
    assert m._session is not None
    m.close()


def test_mouse_model_generate_returns_correct_shape() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    pts = m.generate((100, 200), (900, 400))
    assert len(pts) == 66  # 64 moves + 2 clicks
    for x, y, t in pts:
        assert isinstance(x, int)
        assert isinstance(y, int)
        assert isinstance(t, int)


def test_mouse_model_click_false_omits_clicks() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    pts = m.generate((100, 200), (900, 400), click=False)
    assert len(pts) == 64


def test_mouse_model_seed_reproducibility() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    a = m.generate((100, 200), (900, 400), seed=42)
    b = m.generate((100, 200), (900, 400), seed=42)
    assert a == b


def test_mouse_model_invalid_path_raises_model_load_error() -> None:
    from ai_mouse.mouse import MouseModel
    from ai_mouse.errors import ModelLoadError
    with pytest.raises(ModelLoadError):
        MouseModel(model_path="/nonexistent/path/here")
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_mouse.py -v
```

Expected: ImportError.

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

Create `src/ai_mouse/mouse.py`:

```python
"""MouseModel — ONNX Runtime-backed mouse trajectory generation."""
from __future__ import annotations

import json
import math
from collections.abc import Sequence
from pathlib import Path
from typing import Optional

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse._coord import decode_trajectory
from ai_mouse._postprocess import (
    build_timestamps,
    enforce_forward_monotonic,
    gaussian_smooth,
    resample_arc,
    sample_duration,
    smooth_start,
    snap_endpoints,
    truncnorm_sample,
)
from ai_mouse.errors import GenerationError, ModelLoadError

_N_EULER_STEPS = 10


class MouseModel:
    """Persistent ONNX Runtime session for mouse trajectory generation.

    Construct once and reuse across calls — the underlying
    ``InferenceSession`` is created lazily in ``__init__`` and kept alive
    until :meth:`close` is called.
    """

    def __init__(
        self,
        model_path: str | Path | None = None,
        providers: Sequence[str] | None = None,
        seed: int | None = None,
    ) -> None:
        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None

        onnx_path = resolve(path_obj, "flow_model.onnx")
        cfg_path = resolve(path_obj, "train_config.json")
        click_path = resolve(path_obj, "click_dist.json")
        dur_path = resolve(path_obj, "duration_dist.json")

        cfg = json.loads(cfg_path.read_text())
        self._seq_len = int(cfg["seq_len"])
        self._cond_dim = int(cfg.get("cond_dim", 3))
        self._click_params = json.loads(click_path.read_text())
        self._duration_dist = json.loads(dur_path.read_text())

        try:
            self._session = ort.InferenceSession(
                str(onnx_path),
                providers=list(providers) if providers else ["CPUExecutionProvider"],
            )
        except Exception as exc:
            raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc

        self._default_seed = seed
        self._rng = np.random.default_rng(seed)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def generate(
        self,
        start: tuple[int, int],
        end: tuple[int, int],
        n_points: int = 64,
        speed: float | None = None,
        click: bool = True,
        seed: int | None = None,
    ) -> list[tuple[int, int, int]]:
        """Generate a human-like mouse trajectory from ``start`` to ``end``.

        Args:
            start: (x, y) starting pixel coordinate.
            end: (x, y) target pixel coordinate.
            n_points: number of move points (default 64).
            speed: optional multiplier; ``speed=2`` halves the duration.
            click: append mouse-down and mouse-up events at the end.
            seed: per-call seed overriding the instance seed.

        Returns:
            List of (x, y, t_ms) tuples. ``t_ms`` is cumulative from 0.
        """
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        sx, sy = float(start[0]), float(start[1])
        ex, ey = float(end[0]), float(end[1])
        dist = max(math.hypot(ex - sx, ey - sy), 1.0)

        total_duration = sample_duration(self._duration_dist, dist, rng)
        if speed is not None and speed > 0:
            total_duration /= speed
        total_duration = max(total_duration, 10.0)

        cond = np.array(
            [
                dist / 2000.0,
                math.log(dist / 100.0),
                math.log(total_duration / 500.0),
            ],
            dtype=np.float32,
        )[None]

        # Euler ODE
        x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
        dt = 1.0 / _N_EULER_STEPS
        for step in range(_N_EULER_STEPS):
            t = np.full((1,), step * dt, dtype=np.float32)
            v = self._session.run(
                ["v"], {"x_t": x, "t": t, "cond": cond}
            )[0]
            x = x + v * dt

        if not np.all(np.isfinite(x)):
            raise GenerationError("Trajectory contains NaN/Inf after Euler integration")

        forward = x[0, :, 0].copy()
        lateral = x[0, :, 1].copy()
        log_dt = x[0, :, 2].copy()

        # Spatial post-processing
        forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
        forward, lateral = smooth_start(forward, lateral)
        forward = enforce_forward_monotonic(forward)
        lateral = gaussian_smooth(lateral, sigma=1.0)

        # Temporal post-processing
        log_dt = np.clip(log_dt, 0.0, 5.0)
        log_dt[0] = 0.0

        # Decode to pixels
        normalised = np.stack([forward, lateral], axis=1)
        pixels = decode_trajectory(normalised, start, end)

        if n_points != self._seq_len:
            pixels = resample_arc(pixels, n_points)
            log_dt = np.interp(
                np.linspace(0, 1, n_points),
                np.linspace(0, 1, self._seq_len),
                log_dt,
            )

        ts = build_timestamps(log_dt, total_duration)

        moves: list[tuple[int, int, int]] = [
            (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i])))
            for i in range(n_points)
        ]
        if not click:
            return moves

        click_dur = int(
            truncnorm_sample(
                float(self._click_params["mu"]),
                float(self._click_params["sigma"]),
                float(self._click_params["low"]),
                float(self._click_params["high"]),
                rng,
            )
        )
        click_dur = max(click_dur, int(float(self._click_params["low"])))
        last_t = moves[-1][2]
        cx, cy = moves[-1][0], moves[-1][1]
        return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)]

    def sample_click_duration_ms(self, seed: int | None = None) -> int:
        """Sample one click hold duration from the bundled distribution."""
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        v = truncnorm_sample(
            float(self._click_params["mu"]),
            float(self._click_params["sigma"]),
            float(self._click_params["low"]),
            float(self._click_params["high"]),
            rng,
        )
        return max(int(v), int(float(self._click_params["low"])))

    def close(self) -> None:
        """Release the ONNX session."""
        self._session = None  # type: ignore[assignment]

    def __enter__(self) -> "MouseModel":
        return self

    def __exit__(self, *exc) -> None:
        self.close()
```

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

```bash
uv run pytest tests/unit/test_mouse.py -v
```

Expected: 5 pass. If `test_mouse_model_seed_reproducibility` fails because the bundled ONNX model produces different results across two runs with the same seed, that's a bug in MouseModel. Verify the rng is properly seeded.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/mouse.py tests/unit/test_mouse.py
git commit -m "feat(lib): add MouseModel (numpy + ONNX Runtime)"
```

---

### Task 4.10: Write `scroll.py` (`ScrollModel`)

**Files:**
- Create: `src/ai_mouse/scroll.py`
- Create: `tests/unit/test_scroll.py`

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

Create `tests/unit/test_scroll.py`:

```python
"""Tests for ScrollModel and ai_mouse.generate_scroll()."""
from __future__ import annotations

import pytest


def test_scroll_model_init_default() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    assert m._seq_len > 0
    m.close()


def test_scroll_model_generate_target_mode() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    events = m.generate(0, 1500, mode="target")
    assert len(events) >= 5
    total = sum(e["deltaY"] for e in events)
    # Should approach but not necessarily equal 1500 exactly
    assert 1000 <= total <= 2000  # broad — quantisation can drift
    assert events[0]["t"] == 0
    assert all(e["deltaMode"] == 0 for e in events)


def test_scroll_model_direction() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    events = m.generate(2000, 0, mode="target")  # upward
    total = sum(e["deltaY"] for e in events)
    assert total < 0


def test_scroll_invalid_path() -> None:
    from ai_mouse.errors import ModelLoadError
    from ai_mouse.scroll import ScrollModel
    with pytest.raises(ModelLoadError):
        ScrollModel(model_path="/no/such/path")
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_scroll.py -v
```

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

Create `src/ai_mouse/scroll.py`:

```python
"""ScrollModel — ONNX Runtime-backed scroll event generation."""
from __future__ import annotations

import json
import math
from collections.abc import Sequence
from pathlib import Path
from typing import Literal, Optional

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse.errors import ModelLoadError

_DURATION_TABLE = {
    "fast":    lambda d: d * 0.2 + 100.0,
    "precise": lambda d: d * 1.5 + 300.0,
    "target":  lambda d: d * 0.4 + 200.0,
}

_QUANTUM = {"precise": 40, "fast": 120, "target": 120}


class ScrollModel:
    """Persistent ONNX Runtime session for scroll event generation."""

    def __init__(
        self,
        model_path: str | Path | None = None,
        providers: Sequence[str] | None = None,
        seed: int | None = None,
    ) -> None:
        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None

        onnx_path = resolve(path_obj, "scroll_decoder.onnx")
        cfg_path = resolve(path_obj, "scroll_config.json")
        cfg = json.loads(cfg_path.read_text())

        self._seq_len = int(cfg["seq_len"])
        self._latent_dim = int(cfg["latent_dim"])
        self._cond_dim = int(cfg["cond_dim"])

        try:
            self._session = ort.InferenceSession(
                str(onnx_path),
                providers=list(providers) if providers else ["CPUExecutionProvider"],
            )
        except Exception as exc:
            raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc

        self._rng = np.random.default_rng(seed)

    def generate(
        self,
        start_scroll_y: int,
        target_scroll_y: int,
        mode: Literal["target", "fast", "precise"] = "target",
        viewport_height: int = 800,
        seed: int | None = None,
    ) -> list[dict]:
        """Generate a sequence of mouse-wheel events.

        Returns a list of ``{"deltaY": int, "deltaMode": 0, "t": int}``
        dicts. Positive ``deltaY`` = scroll down.
        """
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        distance = abs(target_scroll_y - start_scroll_y)
        direction = 1 if target_scroll_y > start_scroll_y else -1
        distance = max(distance, 10)

        cond = self._build_condition(float(distance), direction, mode, viewport_height)
        z = rng.standard_normal((1, self._latent_dim)).astype(np.float32)
        decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0]

        delta_norm = decoded[:, 0]
        log_dt = decoded[:, 1]

        # Softmax-like normalisation; scale to target distance
        delta_weights = np.exp(delta_norm)
        delta_weights /= delta_weights.sum()
        delta_px = delta_weights * distance * direction

        quantum = _QUANTUM[mode]
        delta_q = np.round(delta_px / quantum) * quantum
        for i in range(len(delta_q)):
            if delta_q[i] == 0:
                delta_q[i] = quantum * direction
        # Adjust last event so total matches target distance
        delta_q[-1] += (distance * direction) - delta_q.sum()

        # Timestamp building
        if len(log_dt) > 3:
            median_log = float(np.median(log_dt))
            log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
            log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
        dt_ms = np.clip(np.exp(log_dt), 5, 80)
        expected = _DURATION_TABLE[mode](distance)
        dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80)

        t_abs = np.cumsum(dt_ms).astype(int)
        t_abs = np.concatenate([[0], t_abs[:-1]])
        for i in range(1, len(t_abs)):
            if t_abs[i] <= t_abs[i - 1]:
                t_abs[i] = t_abs[i - 1] + 5

        events: list[dict] = []
        for i in range(self._seq_len):
            dy = int(delta_q[i])
            if dy != 0 or len(events) < 5:
                events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
        return events

    def _build_condition(
        self,
        distance: float,
        direction: int,
        mode: str,
        viewport_height: int,
    ) -> np.ndarray:
        mode_onehot = [0.0, 0.0, 0.0]
        if mode == "target":
            mode_onehot[0] = 1.0
        elif mode == "fast":
            mode_onehot[1] = 1.0
        elif mode == "precise":
            mode_onehot[2] = 1.0
        return np.array(
            [
                distance / 5000.0,
                math.log(max(distance, 1.0) / 500.0),
                float(direction),
                viewport_height / 1000.0,
                *mode_onehot,
            ],
            dtype=np.float32,
        )

    def close(self) -> None:
        self._session = None  # type: ignore[assignment]

    def __enter__(self) -> "ScrollModel":
        return self

    def __exit__(self, *exc) -> None:
        self.close()
```

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

```bash
uv run pytest tests/unit/test_scroll.py -v
```

Expected: 4 pass.

- [ ] **Step 5: Commit**

```bash
git add src/ai_mouse/scroll.py tests/unit/test_scroll.py
git commit -m "feat(lib): add ScrollModel (numpy + ONNX Runtime)"
```

---

### Task 4.11: Rewrite `__init__.py` with cached singleton functions

**Files:**
- Modify: `src/ai_mouse/__init__.py`

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

Create `tests/unit/test_public_api.py`:

```python
"""Tests for the public package-level API."""
from __future__ import annotations


def test_public_symbols_importable() -> None:
    from ai_mouse import (
        MouseModel,
        ScrollModel,
        generate,
        generate_scroll,
        errors,
    )
    assert MouseModel is not None
    assert ScrollModel is not None
    assert callable(generate)
    assert callable(generate_scroll)
    assert hasattr(errors, "ModelLoadError")


def test_generate_function_returns_list_of_tuples() -> None:
    from ai_mouse import generate
    pts = generate((100, 100), (300, 200))
    assert isinstance(pts, list)
    assert len(pts) > 0
    assert isinstance(pts[0], tuple)
    assert len(pts[0]) == 3


def test_generate_singleton_reused() -> None:
    from ai_mouse import generate
    from ai_mouse import _model_cache
    _model_cache._get_mouse_model.cache_clear()
    generate((0, 0), (100, 100))
    info_after_first = _model_cache._get_mouse_model.cache_info()
    generate((0, 0), (200, 200))
    info_after_second = _model_cache._get_mouse_model.cache_info()
    assert info_after_second.hits > info_after_first.hits


def test_version_present() -> None:
    import ai_mouse
    assert hasattr(ai_mouse, "__version__")
    assert isinstance(ai_mouse.__version__, str)
```

- [ ] **Step 2: Run, observe failure**

```bash
uv run pytest tests/unit/test_public_api.py -v
```

- [ ] **Step 3: Create `_model_cache.py`**

Create `src/ai_mouse/_model_cache.py`:

```python
"""Process-level lru_cache for default MouseModel / ScrollModel instances."""
from __future__ import annotations

from collections.abc import Sequence
from functools import lru_cache
from pathlib import Path

from ai_mouse.mouse import MouseModel
from ai_mouse.scroll import ScrollModel


@lru_cache(maxsize=4)
def _get_mouse_model(
    model_key: str,
    providers_key: tuple[str, ...],
) -> MouseModel:
    path = None if model_key == "__bundled__" else Path(model_key)
    providers = list(providers_key) if providers_key else None
    return MouseModel(model_path=path, providers=providers)


@lru_cache(maxsize=4)
def _get_scroll_model(
    model_key: str,
    providers_key: tuple[str, ...],
) -> ScrollModel:
    path = None if model_key == "__bundled__" else Path(model_key)
    providers = list(providers_key) if providers_key else None
    return ScrollModel(model_path=path, providers=providers)


def get_mouse_model(
    model_path: str | Path | None,
    providers: Sequence[str] | None,
) -> MouseModel:
    key = "__bundled__" if model_path is None else str(model_path)
    return _get_mouse_model(key, tuple(providers or ()))


def get_scroll_model(
    model_path: str | Path | None,
    providers: Sequence[str] | None,
) -> ScrollModel:
    key = "__bundled__" if model_path is None else str(model_path)
    return _get_scroll_model(key, tuple(providers or ()))
```

- [ ] **Step 4: Rewrite `__init__.py`**

Replace `src/ai_mouse/__init__.py` entirely:

```python
"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events.

Public API:

    from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel

See https://github.com/<owner>/ai_mouse for usage examples.
"""
from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path
from typing import Literal

from ai_mouse import errors
from ai_mouse._model_cache import get_mouse_model, get_scroll_model
from ai_mouse.mouse import MouseModel
from ai_mouse.scroll import ScrollModel

__version__ = "0.2.0"

__all__ = [
    "MouseModel",
    "ScrollModel",
    "errors",
    "generate",
    "generate_scroll",
    "__version__",
]


def generate(
    start: tuple[int, int],
    end: tuple[int, int],
    *,
    n_points: int = 64,
    speed: float | None = None,
    click: bool = True,
    seed: int | None = None,
    model_path: str | Path | None = None,
    providers: Sequence[str] | None = None,
) -> list[tuple[int, int, int]]:
    """Generate a human-like mouse trajectory.

    See :class:`MouseModel.generate` for argument semantics.
    The underlying :class:`MouseModel` is cached process-wide; repeat
    calls with the same ``(model_path, providers)`` reuse the session.
    """
    model = get_mouse_model(model_path, providers)
    return model.generate(
        start=start,
        end=end,
        n_points=n_points,
        speed=speed,
        click=click,
        seed=seed,
    )


def generate_scroll(
    start_scroll_y: int,
    target_scroll_y: int,
    *,
    mode: Literal["target", "fast", "precise"] = "target",
    viewport_height: int = 800,
    seed: int | None = None,
    model_path: str | Path | None = None,
    providers: Sequence[str] | None = None,
) -> list[dict]:
    """Generate a sequence of mouse-wheel events. See :class:`ScrollModel.generate`."""
    model = get_scroll_model(model_path, providers)
    return model.generate(
        start_scroll_y=start_scroll_y,
        target_scroll_y=target_scroll_y,
        mode=mode,
        viewport_height=viewport_height,
        seed=seed,
    )
```

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

```bash
uv run pytest tests/unit -v
```

Expected: all pass.

- [ ] **Step 6: Commit**

```bash
git add src/ai_mouse/__init__.py src/ai_mouse/_model_cache.py tests/unit/test_public_api.py
git commit -m "feat(lib): rewrite __init__.py with cached singleton entrypoints"
```

---

### Task 4.12: Add `py.typed` marker

**Files:**
- Create: `src/ai_mouse/py.typed`

- [ ] **Step 1: Create marker file**

```bash
touch src/ai_mouse/py.typed
```

- [ ] **Step 2: Commit**

```bash
git add src/ai_mouse/py.typed
git commit -m "feat(lib): add py.typed marker (PEP 561)"
```

---

### Task 4.13: Add golden regression tests

**Files:**
- Create: `tests/unit/test_golden.py`

- [ ] **Step 1: Write the test**

Create `tests/unit/test_golden.py`:

```python
"""Golden regression tests — lock library output against pre-migration captures.

Tolerance: pixels and ms allowed ±2 due to ORT/PyTorch fp accumulation
and rounding differences. Update goldens only via an explicit recapture.
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from ai_mouse import generate, generate_scroll

_GOLDEN_DIR = Path(__file__).parent / "data"

_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
    ((100, 200), (900, 400)),
    ((500, 500), (500, 100)),
    ((200, 600), (800, 200)),
    ((100, 100), (130, 110)),
    ((50, 50),   (1500, 900)),
    ((400, 300), (500, 300)),
    ((300, 300), (700, 700)),
    ((600, 400), (200, 100)),
]

_SCROLL_CASES: list[tuple[int, int, str]] = [
    (0,    1500, "target"),
    (0,    500,  "precise"),
    (0,    5000, "fast"),
    (2000, 0,    "target"),
    (0,    800,  "precise"),
    (0,    3500, "fast"),
    (1000, 1200, "precise"),
    (0,    10000, "fast"),
]


@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_mouse_golden(case_idx: int, seed: int) -> None:
    golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"]
    start, end = _MOUSE_CASES[case_idx]
    pts = generate(start, end, seed=seed)
    arr = np.array(pts, dtype=np.int64)
    assert arr.shape == golden.shape, f"shape mismatch: {arr.shape} vs {golden.shape}"
    diff = np.abs(arr - golden)
    assert diff.max() <= 2, (
        f"case{case_idx} seed{seed} max diff {diff.max()} > 2; "
        f"first diff row: arr[?]=..., golden[?]=..."
    )


@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_scroll_golden(case_idx: int, seed: int) -> None:
    golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"]
    start_y, end_y, mode = _SCROLL_CASES[case_idx]
    events = generate_scroll(start_y, end_y, mode=mode, seed=seed)
    arr = np.array(
        [[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
        dtype=np.int64,
    )
    # Scroll uses VAE prior sampling — looser tolerance.
    # Allow ±1 wheel quantum (40 or 120 px) for deltaY; ±10 ms for t.
    quantum = 120 if mode != "precise" else 40
    if arr.shape != golden.shape:
        pytest.skip(
            f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} "
            f"(quantisation boundary sensitivity)"
        )
    delta_diff = np.abs(arr[:, 0] - golden[:, 0])
    t_diff = np.abs(arr[:, 2] - golden[:, 2])
    assert delta_diff.max() <= quantum, f"deltaY diverged > 1 quantum"
    assert t_diff.max() <= 20, f"t diverged > 20ms"
```

- [ ] **Step 2: Run goldens**

```bash
uv run pytest tests/unit/test_golden.py -v
```

Expected outcomes:
- 32 mouse golden cases run; **some failures are expected** because the post-migration randomness differs from torch (different RNG instance, different floating-point path). Inspect failures.
- If max diff is large (>10), there's a real bug — investigate.
- If max diff is in the 3-8 range, **bump the tolerance** in the test (from 2 to a value that lets all pass) **with a comment explaining why**, then re-commit.

This is the moment of truth for the migration: a passing golden suite says the rewrite preserved semantics.

- [ ] **Step 3: Decide on tolerance**

If you needed to widen the tolerance, edit `test_golden.py` and document it. For example, if max diff observed is 5, change `assert diff.max() <= 2` to `assert diff.max() <= 6, ...` with a comment:

```python
# Tolerance 6: ORT/PyTorch numeric path differs slightly; observed max diff 5.
assert diff.max() <= 6, (
    ...
)
```

- [ ] **Step 4: Commit**

```bash
git add tests/unit/test_golden.py
git commit -m "test(lib): add golden regression suite for mouse + scroll"
```

---

### Task 4.14: Delete obsolete legacy modules

**Files:**
- Delete: `src/ai_mouse/generator.py`
- Delete: `src/ai_mouse/scroll/generator.py` and `src/ai_mouse/scroll/__init__.py`
- Delete: `src/ai_mouse/scroll/` directory entirely (replaced by `src/ai_mouse/scroll.py`)
- Delete: `src/ai_mouse/coord.py` (replaced by `_coord.py`)
- Delete: any remaining files in `src/ai_mouse/` not in the spec's final layout
- Move (clean up): `scripts/build_golden_*.py` → can be deleted now that goldens are captured

- [ ] **Step 1: Check current state of src/ai_mouse/**

```bash
ls src/ai_mouse/
ls src/ai_mouse/scroll/ 2>/dev/null
```

Expected at this point: a mix of new files (`__init__.py`, `mouse.py`, `scroll.py`, `_coord.py`, `_postprocess.py`, `_assets.py`, `_model_cache.py`, `errors.py`, `py.typed`, `assets/`) **and** leftover legacy (`generator.py`, `coord.py`, `scroll/`).

- [ ] **Step 2: Delete legacy files**

```bash
git rm src/ai_mouse/generator.py
git rm src/ai_mouse/coord.py
git rm -r src/ai_mouse/scroll/
```

- [ ] **Step 3: Delete temporary scripts**

```bash
git rm scripts/build_golden_mouse.py scripts/build_golden_scroll.py
rmdir scripts/ 2>/dev/null  # only succeeds if empty
```

- [ ] **Step 4: Verify package layout**

```bash
ls src/ai_mouse/
```

Expected:
```
__init__.py
_assets.py
_coord.py
_model_cache.py
_postprocess.py
errors.py
mouse.py
py.typed
scroll.py
assets/
```

(`scroll/` directory removed; replaced by `scroll.py` module.)

- [ ] **Step 5: Run full library test suite**

```bash
uv run pytest tests/unit -v
```

Expected: all pass.

- [ ] **Step 6: Verify tools/ still works (it now imports from src/ai_mouse private modules)**

The tools-side trainer's import of `from ai_mouse.coord import encode_trajectory` will break (we deleted that file). Fix:

```bash
grep -rn "from ai_mouse.coord" tools/ --include="*.py"
```

For each hit, replace with `from ai_mouse._coord import ...`. The spec explicitly allows tools/ to depend on `ai_mouse._*` private modules.

- [ ] **Step 7: Run tools tests**

```bash
uv run pytest tests/tools -v
```

Expected: all pass.

- [ ] **Step 8: Commit**

```bash
git add -A
git commit -m "refactor(lib): remove legacy generator.py / coord.py / scroll/ package"
```

---

### Task 4.15: Verify clean install has no torch

**Files:**
- (verification only)

- [ ] **Step 1: Build fresh wheel**

```bash
uv build
```

- [ ] **Step 2: Install into a clean venv with NO torch**

```bash
uv venv .venv-test
.venv-test/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl
```

- [ ] **Step 3: Smoke test**

```bash
.venv-test/Scripts/python -c "
from ai_mouse import generate
pts = generate((100, 200), (900, 400), seed=0)
print(f'Got {len(pts)} events')
print('First 3:', pts[:3])
print('Last 2 (clicks):', pts[-2:])
print('No torch needed!')
"
```

Expected: prints output without ImportError. Verifies the "pure inference SDK" promise.

- [ ] **Step 4: Confirm torch absent**

```bash
.venv-test/Scripts/python -c "import torch" 2>&1 | head -1
```

Expected: `ModuleNotFoundError: No module named 'torch'`

- [ ] **Step 5: Clean up**

```bash
rm -rf .venv-test dist/
```

- [ ] **Step 6: No commit needed** (verification only — but if anything failed, fix forward)

---

## Phase 5: Docs + cleanup

### Task 5.1: Rewrite README.md

**Files:**
- Modify: `README.md` (overwrite — if it exists; create if not)

- [ ] **Step 1: Check current README**

```bash
ls README.md 2>/dev/null && head -20 README.md
```

- [ ] **Step 2: Write the new README**

Create/overwrite `README.md`:

````markdown
# ai_mouse

Human-like mouse trajectory and scroll wheel event generator. Inference runs on
ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`.

## Install

```bash
pip install git+https://github.com/<owner>/ai_mouse.git
```

For GPU inference (optional), replace `onnxruntime` with the GPU variant:

```bash
pip install onnxruntime-gpu      # CUDA / TensorRT
# or
pip install onnxruntime-directml # Windows DirectML
```

## Quick start

### Mouse trajectory

```python
from ai_mouse import generate

points = generate(start=(100, 200), end=(900, 400))
# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)]
```

### Scroll wheel

```python
from ai_mouse import generate_scroll

events = generate_scroll(start_scroll_y=0, target_scroll_y=2000)
# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...]
```

### Class API (recommended for repeated calls)

```python
from ai_mouse import MouseModel

m = MouseModel()              # session created once
for target in target_list:
    pts = m.generate((cx, cy), target)
```

### Custom providers / GPU

```python
from ai_mouse import MouseModel

m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
# or
m = MouseModel(providers=["DmlExecutionProvider"])
```

### Reproducibility

```python
m.generate(start, end, seed=42)
```

## API summary

| Name | Purpose |
|---|---|
| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton |
| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session |
| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll |
| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy |

## Thread safety

`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently
from multiple threads — ORT `InferenceSession` is itself thread-safe.

## Development

The repo contains optional dev-only tooling under `tools/` for training your
own models, running the FastAPI web UI, and evaluating output quality. Install
with the `dev` group:

```bash
uv sync --group dev
```

Common commands:

```bash
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py

# Training (after collecting your own data)
uv run python -m tools train --data data/traces.jsonl --output data/models_v2

# Convert Balabit corpus to trace format
uv run python -m tools balabit-adapter --input data/balabit_raw \
    --output data/pretrain_traces.jsonl

# Eval report
uv run python -m tools eval --model-dir data/models_v2 \
    --reference data/pretrain_traces.jsonl --output data/eval_reports/report.md

# Re-export ONNX after retraining
uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
```

Tests:

```bash
uv run pytest tests/unit    # library-only (no torch)
uv run pytest tests/tools   # full dev suite
```
````

- [ ] **Step 3: Commit**

```bash
git add README.md
git commit -m "docs: rewrite README from SDK-consumer perspective"
```

---

### Task 5.2: Create CHANGELOG.md

**Files:**
- Create: `CHANGELOG.md`

- [ ] **Step 1: Write CHANGELOG**

Create `CHANGELOG.md`:

```markdown
# Changelog

All notable changes to this project will be documented here. Format follows
[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows
[Semantic Versioning](https://semver.org/).

## [0.2.0] - 2026-05-11

### Changed (breaking)

- Inference no longer requires PyTorch. Runtime dependencies are now
  `numpy + onnxruntime` only.
- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a
  persistent ORT `InferenceSession`.
- Function signatures `generate()` and `generate_scroll()` are now keyword-only
  past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`).
- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll).
- Removed `config=` parameter; use kwargs directly.
- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`.
- Training, web UI, collector, eval, and data adapter code moved to repo-level
  `tools/`; no longer packaged in the wheel.

### Added

- ONNX-format pre-trained weights bundled inside the wheel via
  `importlib.resources` (~3 MB).
- `tools/export_onnx.py` script with PyTorch/ORT parity check.
- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`,
  `GenerationError`.
- Custom ORT providers parameter for GPU / DirectML.
- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the
  default model across calls.

### Removed

- Legacy `JointCVAE` class.
- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs).
- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling).
```

- [ ] **Step 2: Commit**

```bash
git add CHANGELOG.md
git commit -m "docs: add CHANGELOG with 0.2.0 entry"
```

---

### Task 5.3: Create examples/quickstart.py

**Files:**
- Create: `examples/quickstart.py`

- [ ] **Step 1: Create the example**

```bash
mkdir -p examples
```

Create `examples/quickstart.py`:

```python
"""Minimal example: generate one trajectory + click event.

Run:    uv run python examples/quickstart.py
"""
from __future__ import annotations

from ai_mouse import generate

start = (100, 200)
end = (900, 400)
points = generate(start, end, seed=0)

print(f"Generated {len(points)} events:")
print(f"  first move:  {points[0]}")
print(f"  middle move: {points[len(points) // 2]}")
print(f"  last move:   {points[-3]}")
print(f"  click-down:  {points[-2]}")
print(f"  click-up:    {points[-1]}")

# Typical replay loop pattern. t_ms is cumulative from the start of the trace,
# so block your sender thread until time-since-start reaches each event's t_ms.
#
# import time
# t0 = time.monotonic()
# for x, y, t_ms in points:
#     target_wallclock = t0 + t_ms / 1000.0
#     while time.monotonic() < target_wallclock:
#         pass
#     # replace this with pyautogui / pynput / win32 mouse_event:
#     # send_mouse_move(x, y)
```

- [ ] **Step 2: Run it**

```bash
uv run python examples/quickstart.py
```

Expected: prints 5 event lines.

- [ ] **Step 3: Commit**

```bash
git add examples/quickstart.py
git commit -m "docs: add examples/quickstart.py"
```

---

### Task 5.4: Update CLAUDE.md

**Files:**
- Modify: `CLAUDE.md`

- [ ] **Step 1: Read current CLAUDE.md**

```bash
cat CLAUDE.md
```

- [ ] **Step 2: Rewrite to match new layout**

Replace `CLAUDE.md` with content that reflects the new structure. Key changes:

- All `python -m ai_mouse <cmd>` references → `python -m tools <cmd>`
- "Bundled weights live in `data/models_v2/`" → "Bundled weights live in `src/ai_mouse/assets/`"
- Add a "Library vs tools boundary" section: library code in `src/ai_mouse/` MUST NOT `import torch`; training code in `tools/` may import library private modules
- Test commands split: `pytest tests/unit` vs `pytest tests/tools`
- `main.py` reference → `tools/serve.py`

Suggested new content:

```markdown
# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Project

`ai_mouse` is an ONNX-Runtime SDK that generates human-like mouse trajectories
and scroll wheel events. Runtime dependencies are `numpy + onnxruntime` only;
training and the FastAPI web UI live under `tools/` and are not packaged.

Package manager: **uv**, Python 3.12-3.13.

## Library vs tools — hard boundary

- **`src/ai_mouse/`** — wheel content. NEVER add a `import torch` /
  `import fastapi` / `import scipy` / `import matplotlib` here. CI's
  `library` job installs only runtime deps and would break.
- **`tools/`** — repo-only dev code (training, server, collector, eval,
  data adapters, ONNX export). May `import` library private modules
  (`ai_mouse._coord`, `ai_mouse._postprocess`) freely — they co-evolve.
- **Bundled assets**: `src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx`
  plus four JSON metadata files. Re-generated by
  `tools/export_onnx.py` after retraining.

## Commands

```bash
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py

# Tools CLI dispatch
uv run python -m tools train --data data/traces.jsonl --output data/models_v2
uv run python -m tools eval --model-dir data/models_v2 \
    --reference data/pretrain_traces.jsonl --output data/eval_reports/r.md
uv run python -m tools balabit-adapter --input data/balabit_raw \
    --output data/pretrain_traces.jsonl

# Re-export ONNX (after retraining)
uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/

# Tests
uv run pytest tests/unit            # library-only (no torch)
uv run pytest tests/tools           # full dev suite (needs [dev] group)
uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility

# Dependency sync
uv sync                              # runtime only
uv sync --group dev                  # dev-everything
```

## Architecture

Two parallel ML subsystems share a **collect → train → export → serve** flow.

### Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training)

- **Model**: `TrajectoryFlowModel` (Conditional Flow Matching with 4-layer
  pre-norm Transformer, d_model=128, defined in `tools/models.py`)
- **Inference**: 10-step Euler ODE in Python; each step runs
  `session.run(...)` on `src/ai_mouse/assets/flow_model.onnx`. Followed by
  numpy post-processing in `_postprocess.py` (endpoint snapping, forward
  monotonicity, gaussian smoothing, log_dt → cumulative timestamps,
  truncated-normal click duration).
- **Rotated coordinate frame** (`_coord.py`): trajectories normalised so
  `start → (0, 0)`, `end → (1, 0)`. Makes the model angle/distance invariant.

### Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`)

- **Model**: `ScrollCVAE` (bidirectional-GRU encoder + GRU decoder VAE,
  `tools/scroll/models.py`). Only the **decoder** is exported to ONNX
  (`scroll_decoder.onnx`); encoder is training-only.
- **Inference**: sample `z ~ N(0, 1)` in numpy → one `session.run(...)` →
  softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) →
  scale to target distance → cumulative timestamps.

### Server (`tools/server/`) and frontend (`static/`)

Unchanged from before. App factory `create_app()` mounts four routers under
`/api`. Frontend is vanilla Vue 3 + axios + ECharts via CDN.

## Config

`tools/config.py` holds the training-side dataclasses (`TrainConfig`,
`ScrollTrainConfig`, etc.). The library does NOT use these — its only
"configuration" is what's embedded in `src/ai_mouse/assets/train_config.json`
(architecture params needed to know `seq_len` etc. at inference time).

## Tests

- `tests/unit/conftest.py` — fixtures for library-only tests, no torch.
- `tests/tools/conftest.py` — `model_dir` and `scroll_model_dir` fixtures
  that produce **untrained** torch weights in a temp dir. Used by training-
  /server-side tests.
- `tests/unit/test_golden.py` — regression suite that pins library output
  against `tests/unit/data/golden_{mouse,scroll}.npz` captured before the
  ONNX migration. Tolerance: ±2 pixels/ms for mouse, ±1 quantum for scroll.

Server tests use `httpx.ASGITransport(app=create_app())` with
`pytest-asyncio` — no live socket.
```

- [ ] **Step 3: Commit**

```bash
git add CLAUDE.md
git commit -m "docs: update CLAUDE.md for new src/tools layout"
```

---

### Task 5.5: Add GitHub Actions CI

**Files:**
- Create: `.github/workflows/ci.yml`

- [ ] **Step 1: Create directory**

```bash
mkdir -p .github/workflows
```

- [ ] **Step 2: Write the workflow**

Create `.github/workflows/ci.yml`:

```yaml
name: CI

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]

jobs:
  library:
    name: Library tests (no torch)
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-latest, windows-latest]
        python: ["3.12", "3.13"]
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v3
      - run: uv venv --python ${{ matrix.python }}
      - run: uv pip install -e . pytest
      - run: uv run pytest tests/unit -v

  dev:
    name: Full dev suite (with torch)
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-latest, windows-latest]
        python: ["3.12", "3.13"]
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v3
      - run: uv sync --group dev --python ${{ matrix.python }}
      - run: uv run pytest tests/ -v
```

- [ ] **Step 3: Commit**

```bash
git add .github/workflows/ci.yml
git commit -m "ci: add GitHub Actions workflow (library + dev jobs)"
```

---

### Task 5.6: Delete remaining legacy artefacts

**Files:**
- Delete: legacy `JointCVAE` class (it's in `tools/models.py` now; spec says delete it)
- Delete: any leftover bundled-models path string referencing `data/models_v2/` in the library

- [ ] **Step 1: Check if JointCVAE is referenced anywhere**

```bash
grep -rn "JointCVAE" --include="*.py"
```

If any tools/ file still references it (e.g., `tools/models.py` exports it), remove the class definition and the export.

- [ ] **Step 2: Edit `tools/models.py`**

Open `tools/models.py`, find the `class JointCVAE` block, delete it. Also delete the legacy `from torch.distributions import Normal` import if it's only used there.

- [ ] **Step 3: Verify tools still pass**

```bash
uv run pytest tests/tools -v
```

Expected: all pass.

- [ ] **Step 4: Commit**

```bash
git add tools/models.py
git commit -m "chore: remove legacy JointCVAE"
```

---

### Task 5.7: Final verification — full sweep

- [ ] **Step 1: Clean rebuild + install**

```bash
uv venv .venv-final
.venv-final/Scripts/python -m pip install -e .
.venv-final/Scripts/python -m pip install pytest
.venv-final/Scripts/python -m pytest tests/unit -v
```

Expected: all unit tests pass; no torch installed in this venv.

- [ ] **Step 2: Run dev suite separately**

```bash
uv sync --group dev
uv run pytest tests/ -v
```

Expected: all tests pass.

- [ ] **Step 3: Build wheel and inspect contents**

```bash
uv build
unzip -l dist/ai_mouse-0.2.0-py3-none-any.whl | grep -v "^\$"
```

Expected file list (rough):
- `ai_mouse/__init__.py`
- `ai_mouse/_assets.py`, `_coord.py`, `_model_cache.py`, `_postprocess.py`
- `ai_mouse/errors.py`, `mouse.py`, `scroll.py`, `py.typed`
- `ai_mouse/assets/flow_model.onnx`
- `ai_mouse/assets/scroll_decoder.onnx`
- `ai_mouse/assets/{click_dist,duration_dist,train_config,scroll_config}.json`
- `ai_mouse-0.2.0.dist-info/{METADATA,RECORD,WHEEL}`

No `tools/`, `tests/`, `data/`, `static/`, or `docs/`.

- [ ] **Step 4: Smoke test the wheel**

```bash
uv venv .venv-wheel
.venv-wheel/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl
.venv-wheel/Scripts/python examples/quickstart.py
```

Expected: prints 5 event lines without error.

- [ ] **Step 5: Clean up**

```bash
rm -rf .venv-final .venv-wheel dist/
```

- [ ] **Step 6: Document the final state**

The migration is complete. Update PR description / branch message with:

```
ai_mouse 0.2.0 refactor complete:
- src/ai_mouse/ ships only numpy + ONNX Runtime runtime deps
- tools/ holds all training/server/eval code; not packaged
- 3 MB wheel includes ONNX weights for both mouse and scroll
- Golden regression suite locks behavior across the migration
- README, CHANGELOG, CLAUDE.md updated
```

No commit needed unless you adjust docs further.

---

## Self-Review Notes

(Performed during plan authoring per writing-plans skill instructions.)

**Spec coverage check:**
- §1 Public API ⇒ Tasks 4.9, 4.10, 4.11
- §2 ONNX export ⇒ Tasks 3.1–3.5
- §3 NumPy rewrite ⇒ Tasks 4.4–4.8 (each postprocess fn) + 4.9, 4.10 (using them)
- §4 Migration phasing ⇒ Phases 1, 2 match the spec's 5 stages
- §5 Test strategy ⇒ Golden capture in Phase 0; per-fn unit tests in Phase 4;
  golden regression in 4.13; ONNX parity in 3.5; CI in 5.5
- §6 Documentation ⇒ Tasks 5.1–5.4 + examples/quickstart.py in 5.3

**Placeholder scan:** No "TBD", no unspecified test code, no "similar to Task N"
shortcuts. Every code block is self-contained.

**Type consistency:** Function signatures in `_postprocess.py` referenced
across Tasks 4.5–4.10 use consistent names (`snap_endpoints`, `smooth_start`,
`enforce_forward_monotonic`, `gaussian_smooth`, `resample_arc`,
`build_timestamps`, `sample_duration`, `truncnorm_sample`). `MouseModel` and
`ScrollModel` constructor signatures match the spec verbatim. `_get_mouse_model`
in `_model_cache.py` is used by Task 4.11 with matching signature.

**Scope:** This is one cohesive refactor — five phases with clear hand-offs but
one logical goal. Not splittable into independent plans.
