#!/usr/bin/env python3
"""单图 ONNX 推理，输出原图+增强图拼接对比图。

仅依赖 onnxruntime / cv2 / numpy，无 torch 相关操作。
所有参数写在文件顶部，运行时只需指定输入图片。
"""

import time
from pathlib import Path

import cv2
import numpy as np

# ───────────────────────── 默认参数 ─────────────────────────
ONNX_PATH = 'onnx/retinexformer_lol_v1_1x3x224x224.onnx'
MODEL_HEIGHT = 224
MODEL_WIDTH = 224
OUTPUT_DIR = './'
PROVIDER = 'cpu'          # 'cpu' 或 'cuda'
GAP = 8                    # 拼接白缝宽度（像素）
# ────────────────────────────────────────────────────────────

IMAGE_SUFFIXES = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}


def parse_providers(provider):
    if provider == 'cuda':
        return ['CUDAExecutionProvider', 'CPUExecutionProvider']
    return ['CPUExecutionProvider']


def load_rgb_image(path):
    """读取图片，BGR → RGB。"""
    bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
    if bgr is None:
        raise FileNotFoundError(f'无法读取图片: {path}')
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


def preprocess(rgb):
    """RGB uint8 → float32 [0,1] NCHW。"""
    x = rgb.astype(np.float32) / 255.0
    x = np.transpose(x, (2, 0, 1))[None, ...]
    return np.ascontiguousarray(x, dtype=np.float32)


def postprocess(output):
    """模型输出 NCHW float32 → RGB uint8。"""
    if isinstance(output, (list, tuple)):
        output = output[0]
    output = np.asarray(output)
    if output.ndim != 4 or output.shape[0] != 1:
        raise ValueError(f'异常输出 shape: {output.shape}')
    output = np.clip(output[0], 0.0, 1.0)
    rgb = np.transpose(output, (1, 2, 0))
    return (rgb * 255.0 + 0.5).astype(np.uint8)


def run_onnx(onnx_path, input_tensor, providers):
    """ONNX 推理，返回 (输出数组, 耗时秒, 实际 provider 列表)。"""
    import onnxruntime as ort

    session = ort.InferenceSession(onnx_path, providers=providers)
    input_name = session.get_inputs()[0].name
    output_names = [output.name for output in session.get_outputs()]

    start = time.time()
    outputs = session.run(output_names, {input_name: input_tensor})
    elapsed = time.time() - start
    return outputs[0], elapsed, session.get_providers()


def add_label(rgb, label):
    """在图像左上角添加黑色半透明标签。"""
    canvas = rgb.copy()
    bgr = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)
    cv2.rectangle(bgr, (0, 0), (190, 34), (0, 0, 0), thickness=-1)
    cv2.putText(
        bgr,
        label,
        (10, 24),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.75,
        (255, 255, 255),
        2,
        cv2.LINE_AA,
    )
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


def make_side_by_side(original_rgb, enhanced_rgb, gap):
    """左右拼接 Original | Enhanced，中间加白缝，同时标注标签。"""
    if original_rgb.shape[:2] != enhanced_rgb.shape[:2]:
        enhanced_rgb = cv2.resize(
            enhanced_rgb,
            (original_rgb.shape[1], original_rgb.shape[0]),
            interpolation=cv2.INTER_AREA,
        )
    left = add_label(original_rgb, 'Original')
    right = add_label(enhanced_rgb, 'Enhanced')
    if gap <= 0:
        return np.concatenate([left, right], axis=1)
    spacer = np.full((left.shape[0], gap, 3), 255, dtype=np.uint8)
    return np.concatenate([left, spacer, right], axis=1)


def save_rgb_image(path, rgb):
    """保存 RGB 图片为 PNG。"""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(path), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))


def main():
    import argparse
    parser = argparse.ArgumentParser(description='单图 ONNX 推理 → 拼接图')
    parser.add_argument('--input', required=True, help='输入图片路径')
    args = parser.parse_args()

    input_path = Path(args.input)

    # 1. 读原图，记住原始尺寸
    original_rgb = load_rgb_image(input_path)
    orig_h, orig_w = original_rgb.shape[:2]

    # 2. Resize 到模型输入尺寸
    if (orig_h, orig_w) != (MODEL_HEIGHT, MODEL_WIDTH):
        model_rgb = cv2.resize(original_rgb, (MODEL_WIDTH, MODEL_HEIGHT), interpolation=cv2.INTER_AREA)
    else:
        model_rgb = original_rgb

    # 3. 预处理 → 推理
    input_tensor = preprocess(model_rgb)
    onnx_output, elapsed, providers = run_onnx(ONNX_PATH, input_tensor, parse_providers(PROVIDER))

    # 4. 后处理 → 增强图（模型尺寸）
    enhanced_model_size = postprocess(onnx_output)

    # 5. 将增强图还原到原始尺寸
    if (orig_h, orig_w) != (MODEL_HEIGHT, MODEL_WIDTH):
        enhanced_rgb = cv2.resize(enhanced_model_size, (orig_w, orig_h), interpolation=cv2.INTER_AREA)
    else:
        enhanced_rgb = enhanced_model_size

    # 6. 拼接原图 + 增强图
    comparison = make_side_by_side(original_rgb, enhanced_rgb, GAP)

    # 7. 保存
    output_path = Path(OUTPUT_DIR) / f'{input_path.stem}_compare.png'
    save_rgb_image(output_path, comparison)

    print(f'输入:  {input_path}  ({orig_w}x{orig_h})')
    print(f'模型:  {ONNX_PATH}  ({MODEL_WIDTH}x{MODEL_HEIGHT})')
    print(f'输出:  {output_path}  ({comparison.shape[1]}x{comparison.shape[0]})')
    print(f'耗时:  {elapsed:.4f}s')
    print(f'设备:  {providers}')


if __name__ == '__main__':
    main()