import argparse
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort

def hwc_to_chw(img):
    return np.transpose(img, axes=[2, 0, 1]).copy()


def chw_to_hwc(img):
    return np.transpose(img, axes=[1, 2, 0]).copy()


def read_rgb_float(image_path, size=None):
    img = cv2.imread(str(image_path))
    if img is None:
        raise FileNotFoundError(f'Failed to read image: {image_path}')
    img = img[:, :, ::-1]  # BGR -> RGB
    if size is not None:
        img = cv2.resize(img, size, interpolation=cv2.INTER_AREA)
    return img.astype(np.float32) / 255.0


def preprocess(image_path, size=None):
    img = read_rgb_float(image_path, size=size) * 2 - 1
    return hwc_to_chw(img).astype(np.float32)[None, ...]


def get_onnx_input_size(session):
    shape = session.get_inputs()[0].shape  # [N, C, H, W]
    h, w = shape[2], shape[3]
    if isinstance(h, int) and isinstance(w, int):
        return w, h
    return None


def infer_onnx(session, image_path, size=None):
    tensor = preprocess(image_path, size=size)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    out = session.run([output_name], {input_name: tensor})[0]
    out = np.clip(out, -1, 1)
    out = out * 0.5 + 0.5
    return chw_to_hwc(out.squeeze(0))


def load_session(onnx_path, use_cpu=False):
    providers = []
    if not use_cpu and 'CUDAExecutionProvider' in ort.get_available_providers():
        providers.append('CUDAExecutionProvider')
    providers.append('CPUExecutionProvider')
    return ort.InferenceSession(str(onnx_path), providers=providers)


def save_comparison(hazy_path, dehazed, output_path, size=None):
    """拼接原图(hazy)和去雾图(dehazed)，并标注标签"""
    hazy = cv2.imread(str(hazy_path))
    if hazy is None:
        raise FileNotFoundError(f'Failed to read image: {hazy_path}')

    h, w = dehazed.shape[:2]
    if hazy.shape[:2] != (h, w):
        hazy = cv2.resize(hazy, (w, h), interpolation=cv2.INTER_AREA)

    # hazy 是 BGR，dehazed 是 RGB -> 转 BGR 再拼接
    dehazed_bgr = cv2.cvtColor(
        np.round(dehazed * 255).clip(0, 255).astype(np.uint8),
        cv2.COLOR_RGB2BGR,
    )

    pad = 40
    canvas = np.full((h + pad, w * 2, 3), 255, dtype=np.uint8)
    canvas[pad:, :w] = hazy
    canvas[pad:, w:] = dehazed_bgr

    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(canvas, 'Hazy', (w // 2 - 40, pad - 8), font, 0.9, (0, 0, 0), 2, cv2.LINE_AA)
    cv2.putText(canvas, 'Dehazed', (w + w // 2 - 60, pad - 8), font, 0.9, (0, 0, 0), 2, cv2.LINE_AA)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(output_path), canvas)


def main(args):
    onnx_path = Path(args.onnx)
    if not onnx_path.is_file():
        raise FileNotFoundError(f'ONNX model not found: {onnx_path}')

    input_path = Path(args.input)
    if not input_path.is_file():
        raise FileNotFoundError(f'Input image not found: {input_path}')

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    session = load_session(onnx_path, args.cpu)

    onnx_size = get_onnx_input_size(session)
    if args.width > 0 and args.height > 0:
        infer_size = (args.width, args.height)
    elif onnx_size is not None:
        infer_size = onnx_size
    else:
        infer_size = None

    print(f'ONNX:    {onnx_path}')
    print(f'Input:   {input_path}')
    print(f'Size:    {infer_size}')
    print(f'Output:  {output_path}')

    out_img = infer_onnx(session, input_path, size=infer_size)
    save_comparison(input_path, out_img, output_path)
    print('Done.')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Single-image ONNX inference for DehazeFormer.')
    parser.add_argument('--input', required=True, type=str, help='path to input image')
    parser.add_argument('--onnx', required=True, type=str, help='path to ONNX model')
    parser.add_argument('--output', default='output.png', type=str, help='path to output image')
    parser.add_argument('--width', default=-1, type=int, help='resize input width (-1 uses ONNX fixed size)')
    parser.add_argument('--height', default=-1, type=int, help='resize input height (-1 uses ONNX fixed size)')
    parser.add_argument('--cpu', action='store_true', help='force CPU inference')
    main(parser.parse_args())