import os
import torch
import tqdm
import random


import numpy as np
import torch.nn.functional as F

from PIL import Image
from datetime import timedelta
from accelerate import Accelerator, InitProcessGroupKwargs
from dataclasses import dataclass
from torch.utils.data import Dataset, DataLoader, random_split

from diffusers import UNet2DModel, DPMSolverSDEScheduler, DiffusionPipeline


@dataclass
class TrainingConfig:
    image_size = 256  # the generated image resolution
    train_batch_size = 10
    eval_batch_size = 16  # how many images to sample during evaluation
    num_epochs = 400
    gradient_accumulation_steps = 1
    learning_rate = 1e-4
    lr_warmup_steps = 500
    save_image_epochs = 2
    save_model_epochs = 30
    mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = "D:\\trains\\frame_cond_al_clean"  # the model name locally and on the HF Hub

    push_to_hub = False  # whether to upload the saved model to the HF Hub
    hub_private_repo = False
    overwrite_output_dir = True  # overwrite the old model when re-running the notebook
    seed = 0
    num_train_timesteps = 1000


config = TrainingConfig()


data_path = 'D:\\data\\3D-FRONT-SIMPLE-CLEAN'


def is_png_image_ok(file_path):
    try:
        with Image.open(file_path) as img:
            img.verify()
        return True
    except Exception as e:
        print(f"图片损坏或无法读取: {e}")
        return False

class CondDataset(Dataset):
    def __init__(self):
        self.files = []
        # self.cond_imgs = []
        # self.images = []

    def __getitem__(self, index):
        img_size = (config.image_size, config.image_size)
        img_file, cond_file = self.files[index]
        img = Image.open(img_file).resize(img_size).convert('RGB')
        cond_img = Image.open(cond_file).resize(img_size).convert('RGB')
        mirror_type = random.choice(['', 'h', 'v'])
        if mirror_type == 'h':
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            cond_img = cond_img.transpose(Image.FLIP_LEFT_RIGHT)
        elif mirror_type == 'v':
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
            cond_img = cond_img.transpose(Image.FLIP_TOP_BOTTOM)
        return np.array(cond_img), np.array(img)

    def __len__(self):
        return len(self.files)

    def from_path(self, data_path):
        for file in os.listdir(os.path.join(data_path, 'frame')):
            cond_file = os.path.join(data_path, 'frame', file)
            img_file = os.path.join(data_path, 'layout', file)
            if os.path.exists(img_file) and is_png_image_ok(img_file) and is_png_image_ok(cond_file):
                self.files.append((img_file, cond_file))


class MyConditionalUNet2DModel(UNet2DModel):

    def forward(self, x, t, cond):
        net_input = torch.concat((x, cond / 255 * 2 - 1), dim=1).float()
        output = super().forward(net_input, t, return_dict=False)[0]
        return output


class MyDDPMPipeline(DiffusionPipeline):
    model_cpu_offload_seq = "unet"

    def __init__(self, unet, scheduler):
        super().__init__()
        self.register_modules(unet=unet, scheduler=scheduler)

    @torch.no_grad()
    def __call__(self, cond, batch_size=1, generator=None, num_inference_steps=100):
        # Sample gaussian noise to begin loop
        if isinstance(self.unet.config.sample_size, int):
            image_shape = (
                batch_size,
                3,
                self.unet.config.sample_size,
                self.unet.config.sample_size,
            )
        else:
            image_shape = (batch_size, 3, *self.unet.config.sample_size)

        image = torch.randn(image_shape, generator=generator, device=self.device)
        # set step values
        self.scheduler.set_timesteps(num_inference_steps)

        for t in self.progress_bar(self.scheduler.timesteps):
            # 1. predict noise model_output
            model_output = self.unet(image, t, cond)
            # 2. compute previous image: x_t -> x_t-1
            image = self.scheduler.step(model_output, t, image).prev_sample

        # image = (image / 2 + 0.5).clamp(0, 1) * 255
        image = (image.detach().permute(0, 2, 3, 1) / 2 + 0.5).clamp(0, 1) * 255
        # image = image.detach().permute(0, 2, 3, 1).clamp(0, 1) * 255
        image = image.cpu().numpy()
        return image


def make_grid(images, rows, cols):
    b, w, h, _ = images.shape
    grid = Image.new("RGB", size=(cols * w, rows * h))
    for i in range(b):
        image = Image.fromarray(images[i].astype(np.uint8), mode='RGB')
        grid.paste(image, box=(i % cols * w, i // cols * h))
    return grid


def train_and_val(accelerator, model, optimizer, noise_scheduler, train_dataloader, test_dataloader):
    model, optimizer, train_dataloader, test_dataloader = \
        accelerator.prepare(model, optimizer, train_dataloader, test_dataloader)
    device = accelerator.device
    global_step = 0
    # Now you train the model
    for epoch in tqdm.tqdm(range(config.num_epochs)):
        model.train()
        for step, batch in enumerate(train_dataloader):
            cond_images, clean_images = batch
            cond_images = cond_images.permute(0, 3, 1, 2)#.float().to(device)
            clean_images = clean_images.permute(0, 3, 1, 2)#.float().to(device)
            clean_images = clean_images / 255 * 2 - 1
            # Sample noise to add to the images
            noise = torch.randn(clean_images.shape, device=device)
            bs = clean_images.shape[0]

            # Sample a random timestep for each image
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps, (bs,), device=device
            ).long()

            # Add noise to the clean images according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

            with accelerator.accumulate(model):
                # Predict the noise residual
                noise_pred = model(noisy_images, timesteps, cond_images)
                loss = F.mse_loss(noise_pred, noise)
                # loss.backward()
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad()

            logs = {'epoch': epoch, "loss": loss.detach().item(), "step": global_step}
            accelerator.print(logs)
            global_step += 1

        # After each epoch you optionally sample some demo images with evaluate() and save the model

        # pipeline = pipeline.to(device)
        accelerator.wait_for_everyone()
        if accelerator.is_main_process and \
                ((epoch+1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1):
            unet = accelerator.unwrap_model(model)
            pipeline = MyDDPMPipeline(unet, noise_scheduler)
            pipeline.save_pretrained(config.output_dir)

    accelerator.end_training()


if __name__ == '__main__':
    dataset = CondDataset()
    dataset.from_path(data_path)
    if not os.path.exists(config.output_dir):
        os.makedirs(config.output_dir)

    train_dataset, test_dataset = random_split(dataset, [0.9, 0.1])
    train_dataloader, test_dataloader = DataLoader(train_dataset, batch_size=config.train_batch_size, num_workers=16), \
        DataLoader(test_dataset, batch_size=config.eval_batch_size, num_workers=16)
    kwargs = InitProcessGroupKwargs(backend='gloo', timeout=timedelta(seconds=7200))
    accelerator = Accelerator(kwargs_handlers=[kwargs])
    model_path = os.path.join(config.output_dir, 'unet')
    if os.path.exists(model_path):
        print('load model from last checkpoint')
        model = MyConditionalUNet2DModel.from_pretrained(model_path)
    else:
        model = MyConditionalUNet2DModel(sample_size=config.image_size, in_channels=6, out_channels=3,
                                       layers_per_block=2,  # how many ResNet layers to use per UNet block
                                       block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block
                                       down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
                                       up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D","UpBlock2D"))
    noise_scheduler = DPMSolverSDEScheduler(num_train_timesteps=config.num_train_timesteps)
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
    train_and_val(accelerator, model, optimizer, noise_scheduler, train_dataloader, test_dataloader)
