import argparse

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


def main():
    parser = argparse.ArgumentParser(description="ONNX single-image dehazing inference.")
    parser.add_argument("--onnx", default="aodnet_1x3x480x640_sim.onnx", help="ONNX model path")
    parser.add_argument("--input_image", required=True, help="hazy image path")
    parser.add_argument("--output", default="onnx_result.png", help="output path")
    parser.add_argument("--height", type=int, default=480, help="resize height (match ONNX input)")
    parser.add_argument("--width", type=int, default=640, help="resize width (match ONNX input)")
    parser.add_argument("--normalize", action="store_true", help="use (x-0.5)/0.5 normalization")
    args = parser.parse_args()

    img = Image.open(args.input_image).convert("RGB")
    orig_size = img.size
    img = img.resize((args.width, args.height), Image.BILINEAR)

    arr = np.asarray(img).astype(np.float32) / 255.0
    if args.normalize:
        arr = (arr - 0.5) / 0.5
    arr = arr.transpose(2, 0, 1)[None, :, :, :].astype(np.float32)

    sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
    input_name = sess.get_inputs()[0].name
    out = sess.run(None, {input_name: arr})[0]

    out = out[0].transpose(1, 2, 0)
    out = np.clip(out, 0.0, 1.0)
    out_img = Image.fromarray((out * 255.0).round().astype(np.uint8))

    compare = Image.new("RGB", (img.width + out_img.width, max(img.height, out_img.height)))
    compare.paste(img, (0, 0))
    compare.paste(out_img, (img.width, 0))
    compare.save(args.output)

    print("Saved:", args.output)
    print("Input: {:.0f}x{:.0f} -> ONNX: {}x{}".format(*orig_size, args.height, args.width))


if __name__ == "__main__":
    main()