import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, BertModel, BertTokenizer, AutoModelForCausalLM
from torch.cuda.amp import autocast, GradScaler
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model


def create_causal_mask(seq_length):
    """
    创建单向注意力掩码，确保模型只能访问当前和过去的时间步。
    """
    mask = torch.tril(torch.ones(seq_length, seq_length))  # 下三角矩阵
    mask = mask.view(1, seq_length, seq_length)  # 扩展为 [batch_size, seq_length, seq_length]
    return mask

def load_bert_model():
    model_name_or_path = "/root/autodl-tmp/HLLM/plms/bert-base-uncased"
    tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
    model = BertModel.from_pretrained(model_name_or_path)
    return tokenizer, model

def load_llama_model():
    model_name = "/root/autodl-tmp/LLM4KT/plms/TinyLlama_v1.1"
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})  # 添加特殊标记
    model = AutoModel.from_pretrained(model_name)
    return tokenizer, model

def load_llama_lora_model():
    model_name = "/root/autodl-tmp/LLM4KT/plms/TinyLlama_v1.1"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})  # 添加特殊标记
    model = AutoModel.from_pretrained(model_name)
    lora_config = LoraConfig(
        r=32,  # 低秩矩阵的秩
        lora_alpha=32,  # LoRA的缩放因子
        target_modules=["q_proj", "v_proj"],  # 需要应用LoRA的模块
        lora_dropout=0.1
    )
    # 将LoRA配置应用到模型
    model = get_peft_model(model, lora_config)
    return tokenizer, model

def process_text_groups(llm4kt, text_groups):
    outputs = []
    for texts in text_groups:
        output = llm4kt(texts)
        outputs.append(output)
    return outputs


class FeatureAdapter(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim=1024):
        """
        适配器模块，用于对齐特征维度。
        :param input_dim: 输入特征的维度（Llama 的隐藏层维度）
        :param output_dim: 输出特征的维度（BERT 的隐藏层维度）
        :param hidden_dim: 适配器内部的隐藏层维度
        """
        super(FeatureAdapter, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)  # 第一层全连接
        self.fc2 = nn.Linear(hidden_dim, output_dim)  # 第二层全连接
        self.fc3=nn.Linear(input_dim, output_dim)
        self.activation = nn.GELU()  # 使用 GELU 激活函数
        self.dropout = nn.Dropout(0.1)  # 添加 Dropout 防止过拟合

    def forward(self, x):
        """
        前向传播。
        :param x: 输入特征张量，形状为 (batch_size, seq_length, input_dim)
        :return: 对齐后的特征张量，形状为 (batch_size, seq_length, output_dim)
        """
        residual = x  # 保存输入特征作为残差
        x = self.fc1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = x + self.fc3(residual)  # 添加残差连接
        return x

class QueryLLM(nn.Module):
    def __init__(self, model, tokenizer):
        super(QueryLLM, self).__init__()
        self.llm = model
        self.tokenizer = tokenizer
        self.special_token = "[ITEM]"
        self.none_token = "[None]"
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token]})

        self.llm.resize_token_embeddings(len(self.tokenizer))

    def forward(self, texts):
        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=256)
        inputs["input_ids"] = inputs["input_ids"].to(self.llm.device)
        inputs["attention_mask"] = inputs["attention_mask"].to(self.llm.device)
        
        outputs = self.llm(**inputs)
        # 获取特殊标记 [ITEM] 的嵌入
        special_token_id = self.tokenizer.convert_tokens_to_ids(self.special_token)
        special_token_index = (inputs["input_ids"] == special_token_id).nonzero(as_tuple=True)
        query_embeddings = outputs.last_hidden_state[special_token_index]
        return query_embeddings

class ItemLLM(nn.Module):
    def __init__(self, model, tokenizer):
        super(ItemLLM, self).__init__()
        self.llm = model
        self.tokenizer = tokenizer
        self.special_token = "[ITEM]"
        self.none_token = "[None]"
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token]})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        self.scaler = GradScaler()
    def forward(self, texts):
        """
        手动构造输入数据，对每个文本进行编码
        """
        # flat_texts = [text for sublist in texts for text in sublist]
        # # 对展平后的文本进行分词和编码
        # encoded = self.tokenizer.batch_encode_plus(
        #     flat_texts,
        #     add_special_tokens=True,
        #     max_length=64,  # BERT 的最大长度限制
        #     padding='max_length',
        #     truncation=True,
        #     return_attention_mask=True,
        #     return_tensors='pt'
        # )
        # input_ids = encoded['input_ids']
        # attention_mask = encoded['attention_mask']
        # # 将数据移动到设备
        # input_ids = input_ids.to(self.llm.device)
        # attention_mask = attention_mask.to(self.llm.device)
        # 前向传播
        # with torch.no_grad():
        #     outputs = self.llm(input_ids=input_ids, attention_mask=attention_mask)
        """
        自动构造数据，直接传入文本列表
        """
        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=64).to(self.llm.device)
        inputs["input_ids"] = inputs["input_ids"].to(self.llm.device)
        inputs["attention_mask"] = inputs["attention_mask"].to(self.llm.device)
        
        # 获取特殊标记 [ITEM] 的嵌入
        with torch.no_grad():
            outputs = self.llm(**inputs)
        # 获取特殊标记 [ITEM] 的嵌入
        special_token_id = self.tokenizer.convert_tokens_to_ids(self.special_token)
        None_token_id = self.tokenizer.convert_tokens_to_ids(self.none_token)
        
        
        special_token_index = (inputs["input_ids"] == special_token_id).nonzero(as_tuple=True)
        item_embeddings = outputs.last_hidden_state[special_token_index]
        
        none_token_index = (inputs["input_ids"] == None_token_id).nonzero(as_tuple=True)
        none_embeddings = outputs.last_hidden_state[none_token_index]
        
        item_embeddings = torch.cat((item_embeddings, none_embeddings), 0)
        
        
        return item_embeddings

class UserLLM(nn.Module):
    def __init__(self, model):
        super(UserLLM, self).__init__()
        self.llm = model
        self.hidden_size = self.llm.config.hidden_size
        self.prediction_head = nn.Linear(self.hidden_size, self.hidden_size)
        self.causal_mask = create_causal_mask(100)
        
    def forward(self, item_embeddings):
        # bert
        mask = torch.tril(torch.ones(100, 100))  # 下三角矩阵
        mask = mask.view(1, 100, 100)  # 扩展为(batch_size, 
        combined_attention_mask = mask* self.causal_mask
        combined_attention_mask = combined_attention_mask.to(self.llm.device)
        outputs = self.llm(inputs_embeds=item_embeddings, attention_mask=combined_attention_mask)
        user_embeddings = outputs.last_hidden_state[:, :, :]  # 获取用户嵌入
        next_item_embedding = self.prediction_head(user_embeddings)
        return next_item_embedding
        # llama 将项目嵌入作为输入序列
        # input_embeddings = item_embeddings  # (batch_size, seq_len, hidden_size)
        # attention_mask = torch.ones(input_embeddings.size(0), input_embeddings.size(1)).to(item_embeddings.device)
        
        # outputs = self.llm(inputs_embeds=input_embeddings, attention_mask=attention_mask)
        # user_embeddings = outputs.last_hidden_state[:, :, :]  # 获取用户嵌入
        # next_item_embedding = self.prediction_head(user_embeddings)
        # return next_item_embedding

class LLM4KT(nn.Module):
    def __init__(self, user_model):
        super(LLM4KT, self).__init__()
        # self.item_llm = ItemLLM(model=item_model, tokenizer=item_tokenizer)
        self.user_llm = UserLLM(model=user_model)
        
        self.hidden_size=user_model.config.hidden_size
        self.prediction_head = nn.Sequential(
            nn.Linear(2*self.hidden_size, self.hidden_size), 
            nn.ReLU(), 
            nn.Dropout(0.1),
            nn.Linear(self.hidden_size, 256), 
            nn.ReLU(), 
            nn.Dropout(0.1),
            nn.Linear(256, 1)
        )
        self.adapter=FeatureAdapter(768, self.hidden_size)
        self.alignment=FeatureAdapter(self.hidden_size, self.hidden_size)
        self.loss =  nn.BCEWithLogitsLoss(reduction='none')
        self.m = nn.Sigmoid()
        
    def forward(self, input_q, input_qa, target):
        # texts_qa=input_qa
        # outputs = []
        # for texts in texts_qa:
        #     item_embeddings = self.item_llm(texts)
        #     outputs.append(item_embeddings)
        # outputs = torch.stack(outputs)
        # next_item_embedding = self.user_llm(outputs)
        
        # texts_q=input_q
        # outputs_q = []
        # for texts in texts_q:
        #     item_embeddings = self.item_llm(texts)
        #     outputs_q.append(item_embeddings)
        # outputs_q = torch.stack(outputs_q) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        
       
        if self.hidden_size==2048 and input_q.size(-1)==768:
            input_q=self.adapter(input_q)
            input_qa=self.adapter(input_qa)
        else:
            pass
        next_item_embedding = self.user_llm(input_qa)

        outputs_q = input_q
        """"
        将用户嵌入和下一个项目嵌入连接起来，然后通过一个线性层进行预测
        """
        concat_q = torch.cat([next_item_embedding, outputs_q], dim=-1)
        output = self.prediction_head(concat_q)
        labels = target.reshape(-1)
        
        preds = (output.reshape(-1))  # logit
        mask = labels > -0.9
        masked_labels = labels[mask].float()
        masked_preds = preds[mask]
        
        output = self.loss(masked_preds, masked_labels)
        
        return output.sum(), self.m(preds), mask.sum()
        

