"""
Validate the ONNX export against the PyTorch reference.

Three checks:
  1) Random-input parity: torch model logits vs ORT logits (the definitive graph check).
  2) Cactus end-to-end parity: replicate refiners' MVANet eval path (resize 1024 bilinear ->
     image_to_tensor -> ImageNet normalize -> model().sigmoid() -> resize back) for BOTH
     torch and ORT, and compare each to the golden expected_cactus_mask.png.
  3) Save the ORT and torch masks for visual inspection.
"""
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
import torch
from PIL import Image

from refiners.solutions import BoxSegmenter
from refiners.fluxion.utils import image_to_tensor, normalize, tensor_to_image

ROOT = Path(__file__).resolve().parent.parent
MODEL = sys.argv[1] if len(sys.argv) > 1 else str(ROOT / "models" / "mvanet_box_segmenter.onnx")
OUT = ROOT / "assets" / "out"
OUT.mkdir(parents=True, exist_ok=True)


def main() -> None:
    torch.manual_seed(0)
    print(f"[load] torch model + ORT session ({MODEL})")
    seg = BoxSegmenter(device="cpu")
    model = seg.model.eval().float()

    so = ort.SessionOptions()
    sess = ort.InferenceSession(MODEL, sess_options=so, providers=["CPUExecutionProvider"])
    iname = sess.get_inputs()[0].name
    oname = sess.get_outputs()[0].name
    print(f"[ort] input={sess.get_inputs()[0]}  output={sess.get_outputs()[0]}")

    def run_ort(x: torch.Tensor) -> np.ndarray:
        return sess.run([oname], {iname: x.detach().cpu().numpy().astype(np.float32)})[0]

    # 1) random parity
    x = torch.randn(1, 3, 1024, 1024)
    with torch.no_grad():
        yt = model(x).numpy()
    yo = run_ort(x)
    d = np.abs(yt - yo)
    st, so_ = 1 / (1 + np.exp(-yt)), 1 / (1 + np.exp(-yo))
    print("\n[1] RANDOM input parity (logits):")
    print(f"    shapes torch={yt.shape} ort={yo.shape}")
    print(f"    logits  max|d|={d.max():.5e}  mean|d|={d.mean():.5e}")
    print(f"    sigmoid max|d|={np.abs(st - so_).max():.5e}  mean|d|={np.abs(st - so_).mean():.5e}")

    # 2) cactus end-to-end (matches refiners test_mvanet.py: MVANet direct, full image, no box)
    img = Image.open(ROOT / "assets" / "cactus.png").convert("RGB")
    in_t = image_to_tensor(img.resize((1024, 1024), Image.Resampling.BILINEAR)).squeeze()
    in_t = normalize(in_t, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).unsqueeze(0)
    with torch.no_grad():
        pt = model(in_t).sigmoid()
    po = torch.from_numpy(1 / (1 + np.exp(-run_ort(in_t)))).float()
    print("\n[2] CACTUS end-to-end parity (probabilities):")
    print(f"    torch-vs-ort prob max|d|={float((pt - po).abs().max()):.5e} mean|d|={float((pt - po).abs().mean()):.5e}")

    m_t = tensor_to_image(pt).resize(img.size, Image.Resampling.BILINEAR)
    m_o = tensor_to_image(po).resize(img.size, Image.Resampling.BILINEAR)
    exp = Image.open(ROOT / "assets" / "expected_cactus_mask.png").convert("L")
    a_t = np.asarray(m_t).astype(np.int16)
    a_o = np.asarray(m_o).astype(np.int16)
    a_e = np.asarray(exp).astype(np.int16)
    print(f"    MAE(0-255) torch vs expected = {np.abs(a_t - a_e).mean():.3f}")
    print(f"    MAE(0-255) ort   vs expected = {np.abs(a_o - a_e).mean():.3f}")
    print(f"    MAE(0-255) torch vs ort      = {np.abs(a_t - a_o).mean():.3f}")
    m_t.save(OUT / "cactus_torch_mask.png")
    m_o.save(OUT / "cactus_ort_mask.png")
    print(f"    saved {OUT/'cactus_torch_mask.png'} and {OUT/'cactus_ort_mask.png'}")

    # verdict
    ok = float((pt - po).abs().max()) < 5e-2 and np.abs(a_o - a_e).mean() < 3.0
    print("\n[verdict]", "PASS" if ok else "REVIEW",
          "- ONNX matches PyTorch within tolerance" if ok else "- check tolerances above")


if __name__ == "__main__":
    main()