import os

os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([f'{i}' for i in range(8)])


import logging
import random
import torch
import tqdm

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

from PIL import Image
from dataclasses import dataclass
from torch.utils.data import Dataset, DataLoader

from diffusers import UNet2DConditionModel, DPMSolverSDEScheduler, DiffusionPipeline
from transformers import CLIPProcessor, CLIPModel


logger = logging.getLogger()
logger.setLevel(logging.INFO)

file_handler = logging.FileHandler('log.txt')

formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)

logger.addHandler(file_handler)


@dataclass
class TrainingConfig:
    image_size = 256 
    train_batch_size = 6 * 8
    eval_batch_size = 16 
    num_epochs = 40
    gradient_accumulation_steps = 1
    learning_rate = 1e-4
    lr_warmup_steps = 500
    save_image_epochs = 1
    save_model_epochs = 30
    mixed_precision = "fp16" 
    output_dir = "D:\\trains\\text_control_v5" 

    push_to_hub = False 
    hub_private_repo = False
    overwrite_output_dir = True
    seed = 0
    num_train_timesteps = 1000


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 


config = TrainingConfig()


class LayoutDataset(Dataset):
    def __init__(self, processor, max_sentences=50):
        self.files = []
        self.processor = processor
        self.max_sentences = max_sentences
        self.categories = [['bed'], ['sofa'], ['bed', 'sofa'], None]
        self.sample_path = 'D:\\data\\layout_samples'
        if not os.path.exists(self.sample_path):
            os.makedirs(self.sample_path)

    def __getitem__(self, index):
        img_size = (config.image_size, config.image_size)
        img_file, text_file = self.files[index]
        img = Image.open(img_file).resize(img_size).convert('RGB')
        with open(text_file, 'r', encoding='utf8') as f:
            captions = f.readlines()

        if len(captions) < self.max_sentences:
            captions += ['<pad>'] * (self.max_sentences - len(captions))
        elif len(captions) > self.max_sentences:
            captions = random.sample(captions, self.max_sentences)
        text_input_ids = \
        self.processor(text=captions, return_tensors="pt", max_length=10, padding="max_length", truncation=True)[
            'input_ids']
        return torch.from_numpy(np.array(img)), text_input_ids

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

    def from_path(self, data_path):
        img_path = 'D:\\data\\layout_obb_split_v1\\images'
        for x in ['train', 'test', 'val']:
            for file in os.listdir(os.path.join(img_path, x)):
                if file.endswith('.png'):
                    img_file = os.path.join(img_path, x, file)
                    text_file = os.path.join(data_path, 'texts_v2', 'train', file.replace('png', 'txt'))
                    if os.path.exists(text_file) and is_png_image_ok(img_file):
                        self.files.append((img_file, text_file))


class MyConditionalUNet2DModel(UNet2DConditionModel):

    def forward(self, x, t, encoder_hidden_states, **kwargs):
        output = super().forward(x, t, encoder_hidden_states, return_dict=False, **kwargs)[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, encoder_hidden_states, batch_size=1, generator=None, num_inference_steps=100):
        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)
        self.scheduler.set_timesteps(num_inference_steps)

        for t in self.progress_bar(self.scheduler.timesteps):
            model_output = self.unet(image, t, encoder_hidden_states)
            image = self.scheduler.step(model_output, t, image).prev_sample

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


def train_and_val(model, clip_model, optimizer, noise_scheduler, train_dataloader, device, data_num):
    global_step = 0
    tqdm_bar = tqdm.tqdm(total=int(config.num_epochs * data_num / config.train_batch_size))
    for epoch in range(config.num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            clean_images, text_inputs = batch
            clean_images = clean_images.permute(0, 3, 1, 2).to(device)
            clean_images = clean_images / 255 * 2 - 1
            text_inputs = text_inputs.to(device)

            noise = torch.randn(clean_images.shape, device=device)
            bs = clean_images.shape[0]
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps, (bs,), device=device
            ).long()

            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
            text_inputs = torch.reshape(text_inputs, (bs * text_inputs.shape[1], text_inputs.shape[2]))
            encoder_hidden_states = clip_model.get_text_features(input_ids=text_inputs)
            encoder_hidden_states = torch.reshape(encoder_hidden_states, (
            bs, encoder_hidden_states.shape[0] // bs, encoder_hidden_states.shape[1]))
            noise_pred = model(noisy_images, timesteps, encoder_hidden_states)
            loss = F.mse_loss(noise_pred, noise)
            loss.backward()

            optimizer.step()
            optimizer.zero_grad()

            logger.info(f"{epoch}, {global_step}, {loss.detach().item()}")
            tqdm_bar.update(1)
            global_step += 1

        if ((epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1):
            pipeline = MyDDPMPipeline(model.module, noise_scheduler)
            pipeline.save_pretrained(config.output_dir)


if __name__ == '__main__':
    cache_dir = 'D:\\hugging_face'
    data_path = 'D:\\data\\layouts'
    processor = CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14', cache_dir=cache_dir,
                                              local_files_only=True)
    print('finishe load clip')
    dataset = LayoutDataset(processor, max_sentences=3)
    dataset.from_path(data_path)
    data_num = len(dataset)
    print('data num', data_num)
    if not os.path.exists(config.output_dir):
        os.makedirs(config.output_dir)
    train_dataloader = DataLoader(dataset, batch_size=config.train_batch_size, num_workers=4)
    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=3, out_channels=3,
                                         addition_embed_type='text', cross_attention_dim=768,
                                         block_out_channels=(128, 128, 256, 256),
                                         down_block_types=(
                                         "DownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", "DownBlock2D"),
                                         up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D", "UpBlock2D"))
    device = torch.device('cuda')
    model = nn.DataParallel(model.to(device))
    noise_scheduler = DPMSolverSDEScheduler(num_train_timesteps=config.num_train_timesteps)
    clip_model = CLIPModel.from_pretrained('openai/clip-vit-large-patch14', cache_dir=cache_dir,
                                           local_files_only=True)
    clip_model.requires_grad_(False)
    clip_model = clip_model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
    train_and_val(model, clip_model, optimizer, noise_scheduler, train_dataloader, device, data_num)
