"""Export a Moebius UNet checkpoint to the WebGPU-safe ONNX graph Card Anvil loads.

Adapted from simonw/moebius-web python/reexport_unet_no_conv3d.py: the
self-attention pos_conv Conv3d is swapped for the equivalent Conv2d (so the
graph compiles on every WebGPU backend), then exported at a fixed 512 x 512
(64 x 64 latent) with the same input/output names as the hosted unet.onnx.

usage: python export_unet.py <moebius repo> <checkpoint .bin> <out .onnx>
"""
import importlib.util
import os
import sys

import numpy as np
import torch
from torch import nn

repo, weights, out = (os.path.abspath(p) for p in sys.argv[1:4])
sys.path.insert(0, repo)


class PosConv2d(nn.Module):
    """Drop-in for the Conv3d pos_conv; keeps the 5D (b, 1|dim_k, v, hh, ww) layout."""

    def __init__(self, c3: nn.Conv3d):
        super().__init__()
        dim_k = c3.out_channels
        r = c3.kernel_size[1]
        self.conv = nn.Conv2d(1, dim_k, r, padding=r // 2)
        with torch.no_grad():
            self.conv.weight.copy_(c3.weight.squeeze(2))
            self.conv.bias.copy_(c3.bias)

    def forward(self, V):
        b, u, v, hh, ww = V.shape
        x = V.reshape(b * v, 1, hh, ww)
        x = self.conv(x)
        return x.reshape(b, v, -1, hh, ww).permute(0, 2, 1, 3, 4)


class UNetExport(nn.Module):
    def __init__(self, removal_model):
        super().__init__()
        self.m = removal_model

    def forward(self, latent_in9, timesteps, input_ids):
        return self.m(latent_in9, timesteps, input_ids).sample


def load_removal_module():
    # Straight from the file: removal/v1_2/__init__.py drags in the training
    # pipeline and dataset (cv2, transformers, ...), none of which export needs.
    path = os.path.join(repo, "removal", "v1_2", "removal_model.py")
    spec = importlib.util.spec_from_file_location("removal_model", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def main():
    rm = load_removal_module()
    cfg = rm.load_cfg(os.path.join(repo, "config", "model_cfg", "moebius.yaml"))
    model = rm.build_removal_model(cfg, 20).eval()
    state = torch.load(weights, map_location="cpu", weights_only=True)
    print("[load]", model.load_state_dict(state, strict=True))
    print(f"[load] params = {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")

    torch.manual_seed(0)
    latent9 = torch.randn(2, 9, 64, 64)
    ts = torch.tensor([811, 811], dtype=torch.int64)
    ids = torch.tensor([list(range(10, 20)), list(range(10))], dtype=torch.int64)

    wrap = UNetExport(model)
    with torch.no_grad():
        before = wrap(latent9, ts, ids).numpy()
    n = 0
    for mod in model.modules():
        pc = getattr(mod, "pos_conv", None)
        if isinstance(pc, nn.Conv3d):
            mod.pos_conv = PosConv2d(pc)
            n += 1
    with torch.no_grad():
        after = wrap(latent9, ts, ids).numpy()
    d = np.abs(before - after).max()
    print(f"[fix] replaced {n} Conv3d -> Conv2d; torch max|diff| = {d:.3e}")
    assert n > 0 and d < 1e-4, "Conv2d replacement changed the math"

    os.makedirs(os.path.dirname(out), exist_ok=True)
    torch.onnx.export(
        wrap,
        (latent9, ts, ids),
        out,
        opset_version=18,
        dynamo=False,
        input_names=["latent", "timesteps", "input_ids"],
        output_names=["noise"],
        dynamic_axes={
            "latent": {0: "B"},
            "timesteps": {0: "B"},
            "input_ids": {0: "B"},
            "noise": {0: "B"},
        },
    )

    import onnx
    import onnxruntime as ort

    sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"])
    got = sess.run(
        None,
        {"latent": latent9.numpy(), "timesteps": ts.numpy(), "input_ids": ids.numpy()},
    )[0]
    print(f"[check] ONNX vs torch max|diff| = {np.abs(after - got).max():.3e}")
    for i in sess.get_inputs() + sess.get_outputs():
        print(f"[check] {i.name}: {i.type} {i.shape}")
    graph = onnx.load(out, load_external_data=False)
    conv3d = sum(
        1
        for node in graph.graph.node
        if node.op_type == "Conv"
        and any(a.name == "kernel_shape" and len(a.ints) == 3 for a in node.attribute)
    )
    print(f"[check] 3D Conv nodes left: {conv3d}")
    print(f"[done] {out} ({os.path.getsize(out):,} bytes)")


if __name__ == "__main__":
    main()