#!/usr/bin/env python3
"""
Convert an ONNX model (e.g. 2xNomosUni_compact_otf_medium.onnx) to float16,
optimized for browser deployment via onnxruntime-web with the WebGPU
execution provider.

Why fp16 and not int8 for WebGPU:
  Most modern GPUs support FP16 natively -> ~2x size/memory reduction with
  near-equivalent throughput. INT8/INT4 support in WebGPU compute shaders
  is inconsistent across GPU/driver/ORT-web versions, with common gaps in
  fused kernels and integer ops. FP16 is the safe, fast, broadly-supported
  path for WebGPU specifically (as opposed to WASM, where INT8 wins).

What this script does:
  1. Loads and validates the source ONNX model.
  2. Converts weights + compute graph to float16 (keeping a few numerically
     sensitive ops like Resize/Softmax in fp32 via keep_io_types /
     op_block_list, which is standard practice to avoid artifacts).
  3. Verifies the converted model still passes onnx.checker.
  4. Runs a quick numerical sanity check: same input through fp32 and fp16
     models, reports PSNR between the two outputs so you know how close
     the fp16 version is before you ship it.
  5. Saves the result and prints the size comparison.

Usage:
    python convert_fp16_onnx.py \
        --model /Users/emay/Downloads/ONNXmodels/models/2xNomosUni_compact_otf_medium.onnx \
        --out   /Users/emay/Downloads/ONNXmodels/models/2xNomosUni_compact_otf_medium.fp16.onnx
"""

import argparse
import os
import sys

import numpy as np
import onnx

try:
    import onnxruntime as ort
except ImportError:
    print("onnxruntime is required: pip install onnxruntime", file=sys.stderr)
    raise

try:
    from onnxconverter_common import float16
except ImportError:
    print(
        "onnxconverter-common is required: pip install onnxconverter-common",
        file=sys.stderr,
    )
    raise


def psnr(a: np.ndarray, b: np.ndarray, max_val: float = 1.0) -> float:
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    if mse == 0:
        return float("inf")
    return 10 * np.log10((max_val ** 2) / mse)


def inspect(model_path: str) -> onnx.ModelProto:
    print(f"\n{'='*70}\nSTEP 1: Loading & inspecting source model\n{'='*70}")
    model = onnx.load(model_path)
    onnx.checker.check_model(model)

    size_mb = os.path.getsize(model_path) / 1e6
    print(f"File:  {model_path}")
    print(f"Size:  {size_mb:.2f} MB")
    print(f"Input:  {model.graph.input[0].name}")
    print(f"Output: {model.graph.output[0].name}")

    op_counts = {}
    for node in model.graph.node:
        op_counts[node.op_type] = op_counts.get(node.op_type, 0) + 1
    print("Ops:", ", ".join(f"{k}x{v}" for k, v in sorted(op_counts.items())))
    print(f"{'='*70}\n")
    return model


def convert_to_fp16(model: onnx.ModelProto) -> onnx.ModelProto:
    print(f"{'='*70}\nSTEP 2: Converting to float16\n{'='*70}")
    # keep_io_types=True keeps the graph's external input/output tensors as
    # float32 so callers don't need to change how they feed/read data -- ORT
    # inserts Cast nodes at the boundary, cost is negligible vs the conv body.
    # Resize (used for the pixel-shuffle/upsample path in SRVGGNetCompact-like
    # nets) is block-listed since bilinear/nearest resize in fp16 can behave
    # inconsistently across backends; keeping it in fp32 is cheap and safe.
    fp16_model = float16.convert_float_to_float16(
        model,
        keep_io_types=True,
        disable_shape_infer=False,
        op_block_list=["Resize"],
    )
    onnx.checker.check_model(fp16_model)
    print("Conversion complete, model passes onnx.checker.")
    print(f"{'='*70}\n")
    return fp16_model


def sanity_check(orig_model: onnx.ModelProto, fp16_model: onnx.ModelProto, tile: int = 128):
    print(f"{'='*70}\nSTEP 3: Numerical sanity check (fp32 vs fp16 output)\n{'='*70}")

    input_name = orig_model.graph.input[0].name
    x = np.random.rand(1, 3, tile, tile).astype(np.float32)

    sess_orig = ort.InferenceSession(orig_model.SerializeToString(), providers=["CPUExecutionProvider"])
    sess_fp16 = ort.InferenceSession(fp16_model.SerializeToString(), providers=["CPUExecutionProvider"])

    out_orig = sess_orig.run(None, {input_name: x})[0]
    out_fp16 = sess_fp16.run(None, {input_name: x})[0]

    p = psnr(out_orig, out_fp16)
    print(f"PSNR (fp32 vs fp16 output, random calibration tile): {p:.1f} dB")
    if p < 40:
        print("NOTE: PSNR below 40dB -- inspect the fp16 output on a real image "
              "before shipping. This is a synthetic random tile, so also test "
              "with an actual photo for a trustworthy read.")
    else:
        print("Looks good -- fp16 output is numerically very close to fp32.")
    print(f"{'='*70}\n")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True, help="Path to input .onnx model (fp32)")
    ap.add_argument("--out", required=True, help="Path to write the fp16 .onnx model")
    ap.add_argument("--tile", type=int, default=128, help="Tile size for the sanity-check input")
    ap.add_argument("--skip-sanity-check", action="store_true")
    args = ap.parse_args()

    if not os.path.isfile(args.model):
        print(f"Model not found: {args.model}", file=sys.stderr)
        sys.exit(1)

    orig_model = inspect(args.model)
    fp16_model = convert_to_fp16(orig_model)

    if not args.skip_sanity_check:
        sanity_check(orig_model, fp16_model, tile=args.tile)

    onnx.save(fp16_model, args.out)

    orig_size = os.path.getsize(args.model) / 1e6
    new_size = os.path.getsize(args.out) / 1e6
    print(f"Saved: {args.out}")
    print(f"Size:  {orig_size:.2f} MB -> {new_size:.2f} MB "
          f"({(1 - new_size/orig_size)*100:.0f}% smaller)")
    print("\nFor WebGPU in onnxruntime-web, load this model with:\n"
          "  const session = await ort.InferenceSession.create(url, {\n"
          "    executionProviders: ['webgpu']\n"
          "  });\n"
          "Inputs/outputs stay float32 at the JS boundary (keep_io_types=True),\n"
          "so your existing pre/post-processing code doesn't need to change.")


if __name__ == "__main__":
    main()