"""
export_onnx.py — export a trained U-Net to ONNX, auto-detecting channels (v1=2, v2=3).

  python src\\export_onnx.py --ckpt checkpoints\\best.pt     --out web\\radio_unet_v1.onnx
  python src\\export_onnx.py --ckpt checkpoints_v2\\best.pt  --out web\\radio_unet_v2.onnx
"""
import argparse
import numpy as np
import torch

from model import UNet
from dataset import RadioMapSeerDataset


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", default="checkpoints/best.pt")
    p.add_argument("--out", default="web/radio_unet.onnx")
    args = p.parse_args()

    ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
    cargs = ckpt.get("args", {})
    base = cargs.get("base", 64)
    add_dist = bool(cargs.get("distance", False))
    in_ch = 3 if add_dist else 2

    model = UNet(in_channels=in_ch, out_channels=1, base=base)
    model.load_state_dict(ckpt["model_state"])
    model.eval()
    print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist}")

    dummy = torch.randn(1, in_ch, 256, 256)
    torch.onnx.export(model, dummy, args.out,
                      input_names=["input"], output_names=["coverage"],
                      opset_version=17)
    print(f"Exported {args.out}")

    import onnxruntime as ort
    ds = RadioMapSeerDataset(split="test", img_size=256, add_distance=add_dist)
    x, _ = ds[0]
    xb = x.unsqueeze(0).numpy().astype(np.float32)
    with torch.no_grad():
        t_out = model(x.unsqueeze(0)).numpy()
    sess = ort.InferenceSession(args.out, providers=["CPUExecutionProvider"])
    o_out = sess.run(["coverage"], {"input": xb})[0]
    print(f"Max |PyTorch - ONNX|: {np.max(np.abs(t_out - o_out)):.2e}")


if __name__ == "__main__":
    main()