import os
import numpy as np
import cv2
import imageio
import torch
from torchvision.utils import make_grid

def write_image_summary(image_resolution, model, model_input, gt,
                        model_output, writer, total_steps, prefix='train_', save_dtype=None):
    """
    保存预测和GT图像，保持输入图像的通道数和dtype
    """

    # 1. 先恢复成图像 (N C H W) tensor 到 numpy (H W C)
    gt_img = dataio.lin2img(gt['img'], image_resolution)
    pred_img = dataio.lin2img(model_output['model_out'], image_resolution)

    # 都是Tensor: shape (N, C, H, W)
    # 去batch维度，转成HWC
    def tensor_to_hwc(tensor):
        # 假设batch=1
        arr = tensor[0].detach().cpu().numpy()
        # CHW -> HWC
        if arr.ndim == 3:
            arr = np.transpose(arr, (1, 2, 0))
        return arr

    gt_img_np = tensor_to_hwc(gt_img)
    pred_img_np = tensor_to_hwc(pred_img)

    # 2. 根据 save_dtype 转换数据到对应dtype和范围
    def convert_img_to_dtype(img, dtype_target):
        """
        img: np.array, float or int, range unspecified
        dtype_target: numpy dtype or None
        """
        if dtype_target is None:
            # 不转换，直接返回
            return img

        # 若输入是浮点型，先归一化到0~1再映射整型范围
        if np.issubdtype(img.dtype, np.floating):
            img_norm = img - np.min(img)
            max_val = np.max(img_norm)
            if max_val > 0:
                img_norm = img_norm / max_val
            else:
                img_norm = img_norm  # 全零图像
        else:
            # 若本身是整型，先映射到0~1范围
            info = np.iinfo(img.dtype)
            img_norm = img.astype(np.float32) / info.max

        # 如果目标是浮点型，不动。如果是整数型，映射再转换类型
        if np.issubdtype(dtype_target, np.integer):
            dmax = np.iinfo(dtype_target).max
            img_out = (img_norm * dmax).round().astype(dtype_target)
        else:
            img_out = img_norm.astype(dtype_target)

        return img_out

    gt_img_out = convert_img_to_dtype(gt_img_np, save_dtype)
    pred_img_out = convert_img_to_dtype(pred_img_np, save_dtype)

    # 3. 保存路径
    save_dir = os.path.join('./saved_images', prefix.strip('_'))
    os.makedirs(save_dir, exist_ok=True)

    gt_img_path = os.path.join(save_dir, f'{prefix}gt_img_step_{total_steps:06d}.png')
    pred_img_path = os.path.join(save_dir, f'{prefix}pred_img_step_{total_steps:06d}.png')

    # 4. 对于保存，imageio支持各种格式，且保留通道数和dtype
    imageio.imwrite(gt_img_path, gt_img_out)
    imageio.imwrite(pred_img_path, pred_img_out)

    # 5. TensorBoard写入，TensorBoard只能显示uint8，转化为uint8再写
    def to_uint8_for_tb(img):
        if img.dtype == np.uint8:
            return img
        elif np.issubdtype(img.dtype, np.integer):
            info = np.iinfo(img.dtype)
            img_unit8 = (img.astype(np.float32) / info.max * 255).astype(np.uint8)
            return img_unit8
        else:
            # 浮点假定[0,1]
            img_clamped = np.clip(img, 0., 1.)
            return (img_clamped * 255).astype(np.uint8)

    gt_img_tb = to_uint8_for_tb(gt_img_out)
    pred_img_tb = to_uint8_for_tb(pred_img_out)

    # 转成 CHW tensor 写入TensorBoard显示
    writer.add_image(prefix + 'gt_img', torch.from_numpy(np.transpose(gt_img_tb, (2, 0, 1))), global_step=total_steps)
    writer.add_image(prefix + 'pred_img', torch.from_numpy(np.transpose(pred_img_tb, (2, 0, 1))), global_step=total_steps)

    # 6. 继续你的PSNR计算（保持不变）
    write_psnr(dataio.lin2img(model_output['model_out'], image_resolution),
               dataio.lin2img(gt['img'], image_resolution), writer, total_steps, prefix+'img_')
对于这个代码，保存路径应该是什么？