import argparse
import os

import cv2
import numpy as np


def read_image(path, height, width):
    image = cv2.imread(path, cv2.IMREAD_COLOR)
    if image is None:
        raise FileNotFoundError(path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    if image.shape[:2] != (height, width):
        image = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR)
    image = image * 2.0 - 1.0
    image = np.transpose(image, (2, 0, 1))[None].copy()
    return image


def save_comparison(hazy_orig, output, path):
    """拼接原图(hazy)与去雾图(dehazed)并标注文字，还原到原图大小"""
    orig_h, orig_w = hazy_orig.shape[:2]

    # output: NCHW float32 [-1, 1] → 还原到原图大小 HWC uint8 BGR
    output = np.clip(output[0], -1.0, 1.0)
    output = output * 0.5 + 0.5
    output = (np.transpose(output, (1, 2, 0)) * 255).astype(np.uint8)
    if output.shape[:2] != (orig_h, orig_w):
        output = cv2.resize(output, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
    output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)

    h, w = hazy_orig.shape[:2]
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = max(h, w) / 512.0
    thickness = max(1, int(font_scale * 2))
    color = (255, 255, 255)

    # 加 hazy 标注
    cv2.putText(hazy_orig, 'hazy', (int(10 * font_scale), int(30 * font_scale)),
                font, font_scale, color, thickness, cv2.LINE_AA)
    # 加 dehazed 标注
    cv2.putText(output, 'dehazed', (int(10 * font_scale), int(30 * font_scale)),
                font, font_scale, color, thickness, cv2.LINE_AA)

    compare = np.concatenate([hazy_orig, output], axis=1)
    os.makedirs(os.path.dirname(path) or '.', exist_ok=True)
    cv2.imwrite(path, compare)


def main():
    parser = argparse.ArgumentParser(description='Run MixDehazeNet ONNX inference.')
    parser.add_argument('--onnx', default='onnx/MixDehazeNet-s_256x256.onnx', help='Path to simplified ONNX model.')
    parser.add_argument('--image', default='0003_0.8_0.2.jpg', help='Input hazy image.')
    parser.add_argument('--height', type=int, default=256)
    parser.add_argument('--width', type=int, default=256)
    parser.add_argument('--output', default='./onnx_res.png')
    args = parser.parse_args()

    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise SystemExit('Please install ONNX Runtime first: pip install onnxruntime-gpu or onnxruntime') from exc

    # 读取原图用于显示（BGR uint8）
    hazy_orig = cv2.imread(args.image, cv2.IMREAD_COLOR)
    if hazy_orig is None:
        raise FileNotFoundError(args.image)

    input_np = read_image(args.image, args.height, args.width)
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
    session = ort.InferenceSession(args.onnx, providers=providers)
    onnx_output = session.run(['output'], {'input': input_np})[0]
    save_comparison(hazy_orig, onnx_output, args.output)
    print(f'Saved ONNX result: {args.output}')


if __name__ == '__main__':
    main()