human_mouse/examples/cloakbrowser_demo.py

"""End-to-end example of human_mouse driving CloakBrowser.

Subcommands:
    sapi    one real human trajectory, replayed
    multi   N trajectories overlaid on a single canvas in different hues

Usage::

    uv run python examples/cloakbrowser_demo.py sapi
    uv run python examples/cloakbrowser_demo.py multi --n 10

The dataset path defaults to ./sapimouse_data/sapimouse. Override with
--data-root or call human_mouse.download_sapimouse() first.
"""
from __future__ import annotations

import argparse
import json
import math
import statistics
import time
from pathlib import Path

from cloakbrowser import launch

import human_mouse as hm

HERE = Path(__file__).parent
DEMO_HTML = HERE / "demo.html"
DEFAULT_DATA_ROOT = HERE.parent / "sapimouse_data" / "sapimouse"
DEFAULT_OUT = HERE.parent / "outputs"

VIEWPORT = {"width": 1280, "height": 720}
START = (150, 150)
TARGET = (1060, 550)
TARGET_DIST = math.hypot(TARGET[0] - START[0], TARGET[1] - START[1])


def open_page(browser, mode: str):
    context = browser.new_context(viewport=VIEWPORT)
    page = context.new_page()
    page.goto(DEMO_HTML.as_uri() + f"?mode={mode}")
    page.wait_for_selector("#target")
    return page


def save_artifacts(page, name: str, out: Path) -> tuple[Path, Path]:
    out.mkdir(parents=True, exist_ok=True)
    png, js = out / f"{name}.png", out / f"{name}.json"
    js.write_text(json.dumps(page.evaluate("window.__points"), indent=2))
    page.screenshot(path=str(png), full_page=False)
    return png, js


def cmd_sapi(args: argparse.Namespace) -> None:
    print("loading SapiMouse segments...")
    segs = hm.load_all_segments(args.data_root)
    print(f"  -> {len(segs)} segments available")

    seg = hm.pick_segments(segs, n=1,
                           target_distance=TARGET_DIST,
                           max_path_ratio=args.max_ratio,
                           seed=args.seed)[0]
    print(f"  -> {seg.user}/{seg.session}  "
          f"src={len(seg.points)} dist={seg.straight_distance:.0f}px "
          f"dur={seg.duration_ms:.0f}ms ratio={seg.path_ratio:.2f}")

    dense = hm.upsample(hm.affine_warp(seg, START, TARGET), args.density)

    print("launching CloakBrowser (humanize=False)...")
    browser = launch(humanize=False, headless=False)
    page = open_page(browser, "sapi")

    time.sleep(0.3)
    t0 = time.perf_counter()
    hm.replay(page, dense)
    elapsed = (time.perf_counter() - t0) * 1000
    page.locator("#target").click()

    png, js = save_artifacts(page, "sapi", args.out)
    captured = json.loads(js.read_text())
    browser.close()

    print(f"\n  sent={len(dense)} captured={len(captured)} "
          f"replay={elapsed:.0f}ms")
    print(f"  -> {png}")


def cmd_multi(args: argparse.Namespace) -> None:
    multi_dir = args.out / "multi"

    print("loading SapiMouse segments...")
    segs = hm.load_all_segments(args.data_root)
    print(f"  -> {len(segs)} segments total")

    picks = hm.pick_segments(segs, n=args.n,
                             target_distance=TARGET_DIST,
                             max_path_ratio=args.max_ratio,
                             seed=args.seed)

    print("launching CloakBrowser (humanize=False)...")
    browser = launch(humanize=False, headless=False)
    page = open_page(browser, "multi")

    runs = []
    for i, seg in enumerate(picks, 1):
        hue = round(360 * (i - 1) / len(picks))
        page.evaluate(f"window.__setHue({hue}); window.__resetTrace();")
        dense = hm.upsample(hm.affine_warp(seg, START, TARGET), args.density)

        t0 = time.perf_counter()
        hm.replay(page, dense)
        elapsed = (time.perf_counter() - t0) * 1000

        points = page.evaluate("window.__points")
        multi_dir.mkdir(parents=True, exist_ok=True)
        (multi_dir / f"run_{i:02d}.json").write_text(json.dumps(points, indent=2))

        dts = [b["t"] - a["t"] for a, b in zip(points, points[1:])]
        runs.append({
            "i": i, "hue": hue, "user": seg.user, "session": seg.session,
            "src_points": len(seg.points), "sent_points": len(dense),
            "captured_points": len(points),
            "captured_median_dt_ms": round(statistics.median(dts), 2) if dts else 0.0,
            "replay_ms": round(elapsed, 1),
        })
        print(f"  run {i:>2}/{len(picks)} hue={hue:>3} "
              f"{seg.user}/{seg.session} captured={len(points)} pts")

    overlay = multi_dir / "overlay.png"
    page.screenshot(path=str(overlay), full_page=False)
    (multi_dir / "summary.json").write_text(json.dumps({
        "n": len(picks), "runs": runs,
    }, indent=2))
    browser.close()
    print(f"\n  overlay -> {overlay}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="human_mouse demo driven by CloakBrowser",
    )
    parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOT,
                        help="path to unzipped SapiMouse folder")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT,
                        help="output directory for screenshots and JSON")

    sub = parser.add_subparsers(dest="cmd", required=True)

    p_sapi = sub.add_parser("sapi", help="one real SapiMouse trajectory replayed")
    p_sapi.add_argument("--density", type=int, default=4)
    p_sapi.add_argument("--max-ratio", type=float, default=2.0)
    p_sapi.add_argument("--seed", type=int, default=42)
    p_sapi.set_defaults(func=cmd_sapi)

    p_multi = sub.add_parser("multi", help="N trajectories overlaid in different hues")
    p_multi.add_argument("--n", type=int, default=10)
    p_multi.add_argument("--density", type=int, default=4)
    p_multi.add_argument("--max-ratio", type=float, default=2.0)
    p_multi.add_argument("--seed", type=int, default=None)
    p_multi.set_defaults(func=cmd_multi)

    return parser


def main() -> None:
    args = build_parser().parse_args()
    args.func(args)


if __name__ == "__main__":
    main()