#!/usr/bin/env python3
"""
FFDNet ONNX 单图推理脚本
- 读取原图 → 加 AWGN 噪声 → 预处理 → ONNX 推理 → 后处理 → 还原原图大小
- 输出: 原图 + 加噪图 + 结果图 的三图左右拼接（带标注）
- 无 torch 依赖，仅需 cv2 / numpy / onnxruntime
- 所有参数已写为 default 常量，可直接运行
"""

import os
import sys

import cv2
import numpy as np
import onnxruntime as ort

# ============================================================
# 默认参数（可按需修改）
# ============================================================
ONNX_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "model_zoo", "ffdnet_color_fixed_sigma10_640x640_sim.onnx")
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "results", "ffdnet_onnx_infer")
# 若使用非固定 sigma 的模型，设置 USE_FIXED_SIGMA = False 并填写 MODEL_SIGMA
USE_FIXED_SIGMA = True         # True: 单输入 fixed-sigma 模型
MODEL_SIGMA = 10               # 模型 sigma（uint8 尺度，仅 USE_FIXED_SIGMA=False 时生效）
NOISE_SIGMA = 10               # 对原图施加的 AWGN 噪声强度（uint8 尺度，设为 0 则不加噪）
PROVIDERS = ["CPUExecutionProvider"]  # ONNX Runtime 执行后端

# ============================================================


def load_image(img_path):
    """读取图像，返回 BGR uint8 原始图"""
    img_bgr = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
    if img_bgr is None:
        raise FileNotFoundError(f"无法读取图像: {img_path}")
    if img_bgr.ndim == 2:
        img_bgr = cv2.cvtColor(img_bgr, cv2.COLOR_GRAY2BGR)
    return img_bgr


def add_awgn(img_bgr, sigma):
    """对 BGR uint8 图像施加 AWGN，返回加噪后的 BGR uint8 图像"""
    noise = np.random.randn(*img_bgr.shape).astype(np.float32) * sigma
    noisy = img_bgr.astype(np.float32) + noise
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    return noisy


def preprocess(img_bgr, model_h, model_w):
    """
    预处理: BGR->RGB, resize to model size, normalize [0,1], HWC->CHW, add batch
    返回: (1, 3, H, W) float32 张量, 原始尺寸 (h, w)
    """
    orig_h, orig_w = img_bgr.shape[:2]
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, (model_w, model_h), interpolation=cv2.INTER_LINEAR)
    img_float = img_resized.astype(np.float32) / 255.0
    tensor = np.transpose(img_float, (2, 0, 1))[np.newaxis, ...].astype(np.float32)
    return tensor, (orig_h, orig_w)


def postprocess(tensor, orig_h, orig_w):
    """
    后处理: squeeze batch, CHW->HWC, clip [0,1], uint8, resize 回原始尺寸
    返回: BGR uint8 图像
    """
    arr = np.clip(np.squeeze(tensor, axis=0), 0.0, 1.0)          # (3, H, W)
    arr = np.transpose(arr, (1, 2, 0))                            # (H, W, 3)
    arr_uint = (arr * 255.0).round().astype(np.uint8)
    arr_uint = cv2.resize(arr_uint, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
    result_bgr = cv2.cvtColor(arr_uint, cv2.COLOR_RGB2BGR)
    return result_bgr


def make_concat(original_bgr, noisy_bgr, result_bgr,
                label_left="Original", label_mid="Noisy", label_right="FFDNet ONNX"):
    """左中右拼接原图、加噪图与结果图，并添加顶部标注"""
    imgs = [original_bgr, noisy_bgr, result_bgr]
    labels = [label_left, label_mid, label_right]
    hs = [im.shape[0] for im in imgs]
    h = max(hs)

    # 统一高度
    resized = []
    for im in imgs:
        hh, ww = im.shape[:2]
        if hh != h:
            im = cv2.resize(im, (int(ww * h / hh), h), interpolation=cv2.INTER_LINEAR)
        resized.append(im)

    ws = [im.shape[1] for im in resized]
    total_w = sum(ws)
    label_h = max(30, h // 25)
    canvas = np.full((h + label_h, total_w, 3), 255, dtype=np.uint8)

    x = 0
    for idx, im in enumerate(resized):
        canvas[label_h:, x:x + ws[idx]] = im
        x += ws[idx]

    # 标注文字
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = label_h / 30.0
    thickness = max(1, int(font_scale))
    color = (0, 0, 0)

    x = 0
    for idx, (label, w) in enumerate(zip(labels, ws)):
        (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
        cv2.putText(canvas, label, (x + w // 2 - tw // 2, label_h - (label_h - th) // 2),
                    font, font_scale, color, thickness, cv2.LINE_AA)
        x += w

    return canvas


def main():
    if len(sys.argv) < 2:
        print(f"用法: python {os.path.basename(__file__)} <图像路径>")
        print(f"默认 ONNX 模型: {ONNX_PATH}")
        print(f"噪声 sigma: {NOISE_SIGMA}")
        print(f"输出目录: {OUTPUT_DIR}")
        sys.exit(1)

    img_path = sys.argv[1]
    if not os.path.isfile(img_path):
        print(f"错误: 图像不存在: {img_path}")
        sys.exit(1)
    if not os.path.isfile(ONNX_PATH):
        print(f"错误: ONNX 模型不存在: {ONNX_PATH}")
        print(f"请修改脚本顶部 ONNX_PATH 常量，或先导出 ONNX 模型。")
        sys.exit(1)

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # ---- 加载 ONNX 模型 ----
    session = ort.InferenceSession(ONNX_PATH, providers=PROVIDERS)
    input_name = session.get_inputs()[0].name
    _, _, model_h, model_w = session.get_inputs()[0].shape

    # ---- 读取原图 ----
    img_bgr = load_image(img_path)

    # ---- 加噪 ----
    if NOISE_SIGMA > 0:
        noisy_bgr = add_awgn(img_bgr, NOISE_SIGMA)
    else:
        noisy_bgr = img_bgr

    # ---- 预处理（对加噪图） ----
    tensor, (orig_h, orig_w) = preprocess(noisy_bgr, model_h, model_w)

    # ---- ONNX 推理 ----
    if USE_FIXED_SIGMA:
        result = session.run(None, {input_name: tensor})[0]
    else:
        sigma_tensor = np.full((1, 1, 1, 1), MODEL_SIGMA / 255.0, dtype=np.float32)
        sigma_name = session.get_inputs()[1].name
        result = session.run(None, {input_name: tensor, sigma_name: sigma_tensor})[0]

    # ---- 后处理 & 还原原始尺寸 ----
    result_bgr = postprocess(result, orig_h, orig_w)

    # ---- 三图拼接输出 ----
    label_right = f"FFDNet ONNX (sigma={MODEL_SIGMA})" if not USE_FIXED_SIGMA else "FFDNet ONNX"
    label_mid = f"Noisy (sigma={NOISE_SIGMA})" if NOISE_SIGMA > 0 else "Input"
    concat = make_concat(img_bgr, noisy_bgr, result_bgr,
                         label_mid=label_mid, label_right=label_right)

    name = os.path.splitext(os.path.basename(img_path))[0]
    out_path = os.path.join(OUTPUT_DIR, f"{name}_concat.png")
    cv2.imwrite(out_path, concat)
    print(f"输出已保存: {out_path}")


if __name__ == "__main__":
    main()