ai_mouse/tools/eval/__main__.py

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

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

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

import numpy as np

logger = logging.getLogger(__name__)


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

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


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

    ``model_dir`` is accepted for CLI backward compatibility but is no longer
    used — generation goes through the public ai_mouse API which loads the
    bundled ONNX model. Export a fresh .onnx via ``python -m tools.export_onnx``
    to refresh.
    """
    del model_dir  # legacy arg, unused
    from ai_mouse import generate

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


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

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

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

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

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

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

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


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