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

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

import time
from pathlib import Path

import cv2
import numpy as np
import axengine as axe

# ───────────────────────── 默认参数 ─────────────────────────
AXMODEL_PATH = 'Retinexformer_224_224.axmodel'
MODEL_HEIGHT = 224
MODEL_WIDTH = 224
OUTPUT_DIR = './'
GAP = 8                    # 拼接白缝宽度（像素）
# ────────────────────────────────────────────────────────────

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


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 → NCHW uint8。"""
    x = rgb.astype(np.uint8)
    x = np.transpose(x, (2, 0, 1))[None, ...]
    return np.ascontiguousarray(x, dtype=np.uint8)


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 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='单图 axmodel 推理 → 拼接图')
    parser.add_argument('--input', default='1.png', 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)
    session = axe.InferenceSession(AXMODEL_PATH, providers=['AxEngineExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output_names = [output.name for output in session.get_outputs()]

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

    # 4. 后处理 → 增强图（模型尺寸）
    enhanced_model_size = postprocess(axmodel_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) / 'axmodel_res.png'
    save_rgb_image(output_path, comparison)

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


if __name__ == '__main__':
    main()