"""
export_onnx.py — export the trained U-Net to ONNX and verify it matches PyTorch.
Run from the project root:  python src\\export_onnx.py
"""
import numpy as np
import torch

from model import UNet
from dataset import RadioMapSeerDataset

CKPT = "checkpoints/best.pt"
OUT = "web/radio_unet.onnx"

# --- load the trained model ---
ckpt = torch.load(CKPT, map_location="cpu", weights_only=False)
base = ckpt.get("args", {}).get("base", 64)
model = UNet(2, 1, base=base)
model.load_state_dict(ckpt["model_state"])
model.eval()

# --- export to ONNX (fixed 1x2x256x256 input) ---
dummy = torch.randn(1, 2, 256, 256)
torch.onnx.export(
    model, dummy, OUT,
    input_names=["input"], output_names=["coverage"],
    opset_version=17,
)
print(f"Exported {OUT}")

# --- verify ONNX output matches PyTorch on a real test sample ---
import onnxruntime as ort

ds = RadioMapSeerDataset(split="test", img_size=256)
x, _ = ds[0]
xb = x.unsqueeze(0).numpy().astype(np.float32)

with torch.no_grad():
    torch_out = model(x.unsqueeze(0)).numpy()

sess = ort.InferenceSession(OUT, providers=["CPUExecutionProvider"])
onnx_out = sess.run(["coverage"], {"input": xb})[0]

max_diff = float(np.max(np.abs(torch_out - onnx_out)))
print(f"Max |PyTorch - ONNX| on a test sample: {max_diff:.2e}")
print("OK — outputs match" if max_diff < 1e-3 else "WARNING: large mismatch, stop and check")