# 项目创建时间：2024/9/18 18:57
import os

import numpy as np
import torch
import torchvision
from PIL import Image
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.utilities import rank_zero_only


class ImageLogger(Callback):  
    def __init__(self, batch_frequency=2000, max_images=4, clamp=True, increase_log_steps=True,
                 rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False,
                 log_images_kwargs = None,save_dir="./image_log"):
        super().__init__()
        self.rescale = rescale
        self.batch_freq = batch_frequency
        self.max_images = max_images
        if not increase_log_steps:
            self.log_steps = [self.batch_freq]
        self.clamp = clamp
        self.disabled = disabled
        self.log_on_batch_idx = log_on_batch_idx
        self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} # 这里是unconditional_guidance_scale=5.0
        self.log_first_step = log_first_step
        self.save_dir = save_dir

    @rank_zero_only
    def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): # 将生成的图像保存在本地
        root = os.path.join(save_dir, split)
        for k in images:
            grid = torchvision.utils.make_grid(images[k], nrow=4)
            nc = grid.shape[0] # 获取网格图像的通道数。
            if nc > 3: # 如果通道数大于 3，则将图像分成多个 3 通道的子图像，并将它们拼接在一起
                lst = []
                s = 0
                for ch in [3]*(nc//3):
                    lst.append(grid[s:s+ch])
                    s += ch
                grid = torch.cat(lst, dim=1)
            if self.rescale:
                grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
            grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) #(C, H, W)-> (H, W, C)
            # 原来是这样gird = grid.numpy()， 但不支持bfloat16 ->numpy()
            #print(f"grid :{grid.device, type(grid)}")
            #grid = grid.to(torch.float32).cpu().numpy()
            grid = grid.numpy()
            grid = (grid * 255).astype(np.uint8) # 将图像像素值从 [0, 1] 缩放到 [0, 255]
            filename = "{}_gs-{:06}_e-{:06}_b-{:08}.jpg".format(k, global_step, current_epoch, batch_idx)
            path = os.path.join(root, filename)
            os.makedirs(os.path.split(path)[0], exist_ok=True)
            Image.fromarray(grid).save(path)  # 使用 PIL.Image 将 NumPy 数组保存为 JPEG 图像文件。

    def log_img(self, pl_module, batch, batch_idx, split="train"): # 主要功能是在训练过程中记录和保存生成的图像
        check_idx = batch_idx  # if self.log_on_batch_idx else pl_module.global_step
        if (self.check_frequency(check_idx) and  # batch_idx % self.batch_freq == 0
                hasattr(pl_module, "log_images") and
                callable(pl_module.log_images) and
                self.max_images > 0):
            logger = type(pl_module.logger)

            is_train = pl_module.training
            if is_train:
                pl_module.eval()

            with torch.no_grad():
                images = pl_module.log_images(batch, split=split)

            for k in images:
                N = min(images[k].shape[0], self.max_images)
                images[k] = images[k][:N]
                if isinstance(images[k], torch.Tensor):
                    images[k] = images[k].detach().cpu()
                    if self.clamp:
                        images[k] = torch.clamp(images[k], -1., 1.)

            self.log_local(self.save_dir, split, images,
                           pl_module.global_step, pl_module.current_epoch, batch_idx)

            if is_train:
                pl_module.train()

    def check_frequency(self, check_idx):
        return check_idx % self.batch_freq == 0

    # PyTorch Lightning 中的一个回调方法，用于在训练过程中的每个批次结束时执行特定的操作
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): # dataloader_idx
        if not self.disabled:
            self.log_img(pl_module, batch, batch_idx, split="train")
