"""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


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,
    )
    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,
        dynamo=False,
    )
    logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6)
    return out_path