import argparse
import os

import numpy as np
import axengine as axe
from PIL import Image


def preprocess(image_path, height=None, width=None):
    image = Image.open(image_path).convert("RGB")
    if height is not None and width is not None:
        image = image.resize((width, height), Image.BICUBIC)
    # 不做均值归一化，仅转 NCHW/RGB，保存为 uint8
    image_np = np.asarray(image).astype(np.uint8)
    image_np = np.transpose(image_np, (2, 0, 1))[np.newaxis, ...]
    return image_np, image


def postprocess(output):
    output = np.squeeze(output, axis=0)
    output = np.transpose(output, (1, 2, 0))
    output = np.clip(output, 0.0, 1.0)
    return Image.fromarray((output * 255.0).round().astype(np.uint8))


def infer_image(args):
    session = axe.InferenceSession(args.axmodel, providers=["AxEngineExecutionProvider"])
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    input_tensor, hazy_image = preprocess(args.image, args.height, args.width)
    output = session.run([output_name], {input_name: input_tensor})[0]
    dehaze_image = postprocess(output)

    output_dir = os.path.dirname(args.output)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    compare_image = Image.new("RGB", (hazy_image.width + dehaze_image.width, hazy_image.height))
    compare_image.paste(hazy_image, (0, 0))
    compare_image.paste(dehaze_image, (hazy_image.width, 0))
    compare_image.save(args.output)
    print("Comparison image saved to:", args.output)


def parse_args():
    parser = argparse.ArgumentParser(description="Run Light-DehazeNet axmodel inference on one image.")
    parser.add_argument("-m", "--axmodel", default="./LightDehazeNet.axmodel", help="path to axmodel model")
    parser.add_argument("-i", "--image", default='query_hazy_images/outdoor_natural/nh(5).png', help="path to input hazy image")
    parser.add_argument("-o", "--output", default="axmodel_dehaze.jpg", help="path to save side-by-side comparison image")
    parser.add_argument("--height", type=int, default=640, help="resize input to this height before inference")
    parser.add_argument("--width", type=int, default=480, help="resize input to this width before inference")
    args = parser.parse_args()

    if (args.height is None) != (args.width is None):
        parser.error("--height and --width must be specified together")
    return args


if __name__ == "__main__":
    infer_image(parse_args())