import os
import torch
import torchvision.utils as vutils
from einops import rearrange
from torchvision.transforms.functional import to_pil_image


import numpy as np
from PIL import Image
import cv2


def save_output_images(
    output_color: torch.Tensor,
    save_dir: str = "./output_images",
    prefix: str = "output"
):
    """
    保存 output.color 张量中的所有图像为 PNG 文件。

    Args:
        output_color (Tensor): shape = [b, v, c, h, w]
        save_dir (str): 保存图片的目录
        prefix (str): 文件名前缀
    """
    os.makedirs(save_dir, exist_ok=True)

    # 将 [b, v, c, h, w] 变为 [(b * v), c, h, w]
    images = rearrange(output_color, "b v c h w -> (b v) c h w")

    # 遍历每张图像保存
    for i, img_tensor in enumerate(images):
        # 确保数值范围在 [0, 1] 之间
        img_tensor = img_tensor.clamp(0, 1)

        # 转为 PIL 图像，需调整通道顺序为 HWC，并且转换为 uint8 类型
        img_tensor = img_tensor.permute(1, 2, 0)  # 从 (C, H, W) 转为 (H, W, C)
        # 使用 .detach() 来切断计算图，防止在转换时出现问题
        img_pil = Image.fromarray((img_tensor.detach().cpu().numpy() * 255).astype("uint8"))
        # 保存图像
        img_pil.save(os.path.join(save_dir, f"{prefix}_{i:04d}.png"))


def save_depth_images(depth_tensor: torch.Tensor,
                      out_dir: str,
                      prefix: str = "depth",
                      cmap: int = cv2.COLORMAP_INFERNO):
    """
    保存深度图为伪彩色 PNG。

    Args:
        depth_tensor (torch.Tensor): shape [B, V, H, W], float32
        out_dir (str): 保存目录，会自动创建
        prefix (str): 文件名前缀, 最终文件名会是 prefix_b{b}_v{v}.png
        cmap (int): OpenCV 调色板，如 COLORMAP_INFERNO
    """
    os.makedirs(out_dir, exist_ok=True)
    depth_np = depth_tensor.detach().cpu().numpy()
    B, V, H, W = depth_np.shape

    for b in range(B):
        for v in range(V):
            d = depth_np[b, v]
            # 归一化到 [0,255]
            d_min, d_max = float(d.min()), float(d.max())
            if d_max > d_min:
                d_norm = (d - d_min) / (d_max - d_min)
            else:
                d_norm = np.zeros_like(d)
            d_8u = (d_norm * 255).astype(np.uint8)

            # 伪彩色
            colored = cv2.applyColorMap(d_8u, cmap)[:, :, ::-1]  # BGR->RGB

            # 保存
            fname = f"{prefix}_b{b}_v{v}.png"
            path = os.path.join(out_dir, fname)
            Image.fromarray(colored).save(path)

    print(f"Saved {B*V} depth images to {out_dir}")