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

Usage:
    uv run python -m tools.export_onnx \
        --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


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)


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


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)


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())