import os
import torch
import torch.nn as nn
import torch.optim as optim
import pytorch_lightning as pl
from torch.utils.data import DataLoader
from torchvision import transforms
from datasets import load_dataset
from transformers import CLIPProcessor, CLIPModel, ViTModel, GPT2LMHeadModel, GPT2Tokenizer
from torch.nn import functional as F
# 定义生成器
class Vit_Gpt2(nn.Module):
    def __init__(self):
        super(Vit_Gpt2, self).__init__()
        self.vit = ViTModel.from_pretrained("google/vit-base-patch16-224")
        self.gpt2 = GPT2LMHeadModel.from_pretrained("gpt2")
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        # 添加一个线性层来连接ViT和GPT2
        #self.linear = nn.Linear(self.vit.config.hidden_size, self.gpt2.config.n_embd)
        self.mapper = MultiheadDiffAttn(embed_dim=768, depth=4, num_heads=12)

    def forward(self, pixel_values):
        vit_outputs = self.vit(pixel_values=pixel_values).last_hidden_state
        #vit_pooled = vit_outputs.mean(dim=1)  # 对序列维度取平均
        gpt_input = self.mapper(vit_outputs)
        
        gpt_outputs = self.gpt2(inputs_embeds=gpt_input).logits
        
        return gpt_outputs
        
#         tokens = gpt_outputs.argmax(dim=-1)
#         captions = self.tokenizer.batch_decode(tokens)
#         return captions

# 定义判别器
class CLIPDiscriminator(nn.Module):
    def __init__(self):
        super(CLIPDiscriminator, self).__init__()
        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    def forward(self, images, texts):
        # 处理文本输入
        inputs = self.processor(
            text=texts, 
            return_tensors="pt", 
            padding=True, 
            truncation=True
        ).to(images.device)
        
        # 获取文本特征
        text_features = self.clip_model.get_text_features(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask
        )
        
        # 获取图像特征
        image_features = self.clip_model.get_image_features(pixel_values=images)
        
        # 归一化特征向量
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        
        # 计算余弦相似度
        # logits_per_image = (image_features @ text_features.t()) * self.clip_model.logit_scale.exp()
        similarity = torch.sum(image_features * text_features, dim=-1)
        
        return similarity


class GAN(pl.LightningModule):
    def __init__(self, lr=1e-4):
        super(GAN, self).__init__()
        self.generator = Vit_Gpt2()
        self.discriminator = CLIPDiscriminator()
        self.lr = lr
        self.automatic_optimization = False
        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.scaler = torch.cuda.amp.GradScaler()
        
        # 只训练 difftransformer
        for param in [*self.generator.vit.parameters(), *self.discriminator.parameters()]:
            param.requires_grad = False
        
        for param in [*self.generator.vit.pooler.dense.parameters(), 
                     *self.discriminator.clip_model.visual_projection.parameters(), 
                     *self.discriminator.clip_model.text_projection.parameters(),
                     *self.generator.gpt2.parameters()
                     ]:
            param.requires_grad = True

        
    def training_step(self, batch, batch_idx):
        opt_g, opt_d = self.optimizers()
        images, captions = batch['image'], batch['caption'][0]

        # 训练生成器
        opt_g.zero_grad()
        
        # 1. 获取ViT特征并通过mapper转换
        vit_outputs = self.generator.vit(pixel_values=images).last_hidden_state
        gpt_input = self.generator.mapper(vit_outputs)
        
        # 2. 准备语言模型训练数据
        inputs = self.tokenizer(captions, 
                                padding='max_length', 
                                truncation=True, 
                                return_tensors="pt",
                                max_length=197).to(self.device)
        
        input_ids = inputs['input_ids']
        
        # 3. 计算语言模型损失
        outputs = self.generator.gpt2(inputs_embeds=gpt_input, labels=input_ids)
        g_lm_loss = outputs.loss
        
        # 4. 使用generate方法生成文本
        generated_ids = self.generator.gpt2.generate(
            inputs_embeds=gpt_input,
            max_new_tokens=197,  # 生成的文本长度
            num_return_sequences=1,
            do_sample=True
        )
        generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

        # 5. 计算对抗损失
        fake_scores = self.discriminator(images, generated_text)
        g_adv_loss = -fake_scores.mean()
        
        # 6. 计算总损失
        lambda_lm = 1.0
        lambda_adv = 0.1
        g_loss = lambda_lm * g_lm_loss + lambda_adv * g_adv_loss

        # 检查损失值是否有效
        if not torch.isfinite(g_loss):
            print(f"Warning: g_loss is {g_loss}")
            return None

        # 7. 生成器反向传播
        self.manual_backward(g_loss)
        torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=1.0)
        opt_g.step()
        opt_g.zero_grad()

        # 训练判别器
        opt_d.zero_grad()
        
        # 8. 计算真实样本和生成样本的得分
        real_scores = self.discriminator(images, captions)
        fake_scores = self.discriminator(images, generated_text)
        
        # 9. 计算判别器损失
        d_loss = F.relu(1.0 - real_scores).mean() + F.relu(1.0 + fake_scores).mean()

        # 检查损失值是否有效
        if not torch.isfinite(d_loss):
            print(f"Warning: d_loss is {d_loss}")
            return None

        # 10. 判别器反向传播
        self.manual_backward(d_loss)
        torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=1.0)
        opt_d.step()
        opt_d.zero_grad()

        # 11. 记录损失
        self.log('g_loss', g_loss, prog_bar=True)
        self.log('d_loss', d_loss, prog_bar=True)
        self.log('g_lm_loss', g_lm_loss, prog_bar=True)
        self.log('g_adv_loss', g_adv_loss, prog_bar=True)

        return {'loss': g_loss + d_loss}



    def configure_optimizers(self):
        opt_g = torch.optim.Adam(filter(lambda p: p.requires_grad, self.generator.parameters()), lr=self.lr)
        opt_d = torch.optim.Adam(filter(lambda p: p.requires_grad, self.discriminator.parameters()), lr=self.lr)
        return [opt_d, opt_g]





# 数据加载
def get_data_loaders(batch_size=32):
    # 加载Flickr30k数据集
    dataset = load_dataset("nlphuji/flickr30k", cache_dir='/root/autodl-tmp/.autodl/huggingface/datasets/nlphuji___flickr30k')
    
    # 定义图像预处理
    image_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # 定义数据集预处理函数
    def preprocess_function(examples):
        images = [image_transform(img.convert("RGB")) for img in examples["image"]]
        return {"image": images, "caption": examples["caption"]}

    # 应用预处理
    dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset["test"].column_names)
    
    # 设置数据格式
    dataset.set_format(type="torch", columns=["image", "caption"])

    # 创建数据加载器
    train_loader = DataLoader(dataset["test"], batch_size=batch_size, shuffle=True)
    
    return train_loader


检查代码逻辑是否有问题