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 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.query_token = "[QUERY]"
        self.pad_token = "[None1]"
        
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token,self.query_token,self.pad_token]})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        self.scaler = GradScaler()

        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] 的嵌入
        None_token_id = self.tokenizer.convert_tokens_to_ids(self.pad_token)
        none_token_index = (inputs["input_ids"] == None_token_id).nonzero(as_tuple=True)
        none_embeddings = outputs.last_hidden_state[none_token_index] 
        
        query_token_id = self.tokenizer.convert_tokens_to_ids(self.query_token)
        query_token_index = (inputs["input_ids"] == query_token_id).nonzero(as_tuple=True)
        query_embeddings = outputs.last_hidden_state[query_token_index] 
        
        query_embeddings = torch.cat((query_embeddings, none_embeddings), 0)
        
        return query_embeddings

class InterLLM(nn.Module):
    def __init__(self, model, tokenizer):
        super(InterLLM, self).__init__()
        self.llm = model
        self.tokenizer = tokenizer
        self.special_token = "[ITEM]"
        self.none_token = "[None]"
        
        self.query_token = "[QUERY]"
        self.pad_token = "[None1]"
        
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token,self.query_token,self.pad_token]})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        self.scaler = GradScaler()

    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]
             
        # pad_token_id = self.tokenizer.convert_tokens_to_ids(self.pad_token)
        
        None_token_id = self.tokenizer.convert_tokens_to_ids(self.none_token)
        none_token_index = (inputs["input_ids"] == None_token_id).nonzero(as_tuple=True)
        none_embeddings = outputs.last_hidden_state[none_token_index] 
        
        query_embeddings = torch.cat((query_embeddings, none_embeddings), 0)
        
        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.query_token = "[QUERY]"
        self.pad_token = "[None1]"
        
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token,self.query_token,self.pad_token]})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        self.scaler = GradScaler()
    def forward(self, texts):
        """
        自动构造数据，直接传入文本列表
        """
        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=256).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] 的嵌入
   
        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)
        
        # 获取特殊标记 [QUERY] 的嵌入
        query_token_id = self.tokenizer.convert_tokens_to_ids(self.query_token)
        # pad_token_id = self.tokenizer.convert_tokens_to_ids(self.pad_token)
         
        query_token_index = (inputs["input_ids"] == query_token_id).nonzero(as_tuple=True)
        query_embeddings = outputs.last_hidden_state[query_token_index]
        
        # pad_token_index = (inputs["input_ids"] == pad_token_id).nonzero(as_tuple=True)
        # pad_embeddings = outputs.last_hidden_state[pad_token_index]
        
        query_embeddings = torch.cat((query_embeddings, none_embeddings), 0)
       
        
        return item_embeddings, query_embeddings
    
    
    


class ItemLLM_M(nn.Module):
    def __init__(self, model, tokenizer):
        super(ItemLLM_M, self).__init__()
        self.llm = model
        self.tokenizer = tokenizer
        self.special_token = "[ITEM]"
        self.none_token = "[None]"
        
        self.query_token = "[QUERY]"
        self.pad_token = "[None1]"
        
        self.meo_token = "[MEO]"
        self.pad_token2 = "[None2]"
        
        self.tokenizer.add_special_tokens({"additional_special_tokens": [self.special_token, self.none_token,self.query_token,self.pad_token,self.meo_token,self.pad_token2]})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        self.scaler = GradScaler()
    def forward(self, texts):
        """
        自动构造数据，直接传入文本列表
        """
        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=256).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] 的嵌入
   
        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)
        
        # 获取特殊标记 [QUERY] 的嵌入
        query_token_id = self.tokenizer.convert_tokens_to_ids(self.query_token)
        # pad_token_id = self.tokenizer.convert_tokens_to_ids(self.pad_token)
         
        query_token_index = (inputs["input_ids"] == query_token_id).nonzero(as_tuple=True)
        query_embeddings = outputs.last_hidden_state[query_token_index]
        
        # pad_token_index = (inputs["input_ids"] == pad_token_id).nonzero(as_tuple=True)
        # pad_embeddings = outputs.last_hidden_state[pad_token_index]
        
        query_embeddings = torch.cat((query_embeddings, none_embeddings), 0)
        
        
        # 获取特殊标记 [MEO] 的嵌入
        meo_token_id = self.tokenizer.convert_tokens_to_ids(self.meo_token)
        # pad_token_id = self.tokenizer.convert_tokens_to_ids(self.pad_token)
         
        meo_token_index = (inputs["input_ids"] == meo_token_id).nonzero(as_tuple=True)
        meo_embeddings = outputs.last_hidden_state[meo_token_index]
        
        # pad_token_index = (inputs["input_ids"] == pad_token_id).nonzero(as_tuple=True)
        # pad_embeddings = outputs.last_hidden_state[pad_token_index]
        
        meo_embeddings = torch.cat((meo_embeddings, none_embeddings), 0)
       
        
        return item_embeddings, query_embeddings,meo_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 ,label_matrix):

        #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)
        # 假设 label 矩阵变量名为 label_matrix，其形状为 (batch_size, seq_len, 1)
        attention_mask = (label_matrix.squeeze(-1) != -1).float().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_ID(nn.Module):
    def __init__(self,query_model,query_tokenizer, item_model,item_tokenizer, user_model):
        super(LLM4KT_ID, self).__init__()
        self.query_llm = QueryLLM(model=query_model, tokenizer=query_tokenizer)
        self.item_llm = InterLLM(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(self.hidden_size, 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):

        #query and item
        
        
        # merged_texts = [g1 + g2 for g1,g2 in zip(input_q, input_qa)]
        # merged_texts = input_qa
        
        # outputs_qa = []
        # outputs_q=[]
        # for texts in merged_texts:
        #     item_embeddings,query_embeddings = self.item_llm(texts)
        #     outputs_qa.append(item_embeddings)
        #     outputs_q.append(query_embeddings)
        # outputs_q = torch.stack(outputs_q) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        # outputs_qa = torch.stack(outputs_qa) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        
        #no need to use query
        
        
        outputs_qa = []
        outputs_q=[]
        texts_q=input_q
        for texts in texts_q:
            # item_embeddings = self.item_llm(texts)
            item_embeddings = self.query_llm(texts)
            outputs_q.append(item_embeddings)
        outputs_q = torch.stack(outputs_q) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        
        texts_qa=input_qa
        outputs_qa = []
        for texts in texts_qa:
            # item_embeddings = self.item_llm(texts)
            item_embeddings = self.item_llm(texts)
            outputs_qa.append(item_embeddings)
        outputs_qa = torch.stack(outputs_qa) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        
       
        
        input_q=self.adapter(outputs_q)
        input_qa=self.adapter(outputs_qa)
        
        next_item_embedding = self.user_llm(input_qa, target)

        """"
        将用户嵌入和下一个项目嵌入连接起来，然后通过一个线性层进行预测
        """
        # concat_q = torch.cat([next_item_embedding, outputs_q], dim=-1)
        concat_qa = torch.cat([next_item_embedding, input_q], dim=-1)
        output = self.prediction_head(concat_qa)
        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()
    


# 该模型使用了 LLM4KT_ID 模型，但是分为三部分，分别对query和inter和user进行建模
class LLM4KT_ID_IU(nn.Module):
    def __init__(self,item_model,item_tokenizer, user_model):
        super(LLM4KT_ID_IU, self).__init__()
        
        # self.item_llm = ItemLLM(model=item_model, tokenizer=item_tokenizer)
        self.item_llm = ItemLLM_M(model=item_model, tokenizer=item_tokenizer)
        self.user_llm = UserLLM(model=user_model)
        
        
        self.input_size=item_model.config.hidden_size
        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.meomry_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(self.input_size, self.hidden_size)

        self.adapter1=FeatureAdapter(self.input_size, self.hidden_size)        
        
        self.alignment=FeatureAdapter(self.input_size, self.hidden_size)
        
        
        self.loss =  nn.BCEWithLogitsLoss(reduction='none')
        self.m = nn.Sigmoid()
        
    def forward(self, input_q, input_qa,input_meo, target,meo_target):

        #query and item

        merged_texts = [g1 + g2 +g3 for g1,g2,g3 in zip(input_q, input_qa,input_meo)]
        # merged_texts = input_qa
        
        outputs_qa = []
        outputs_q=[]
        outputs_m=[]
        for texts in merged_texts:
            item_embeddings,query_embeddings,meo_embeddings = self.item_llm(texts)
            outputs_qa.append(item_embeddings)
            outputs_q.append(query_embeddings)
            outputs_m.append(meo_embeddings)
        outputs_q = torch.stack(outputs_q) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        outputs_qa = torch.stack(outputs_qa) # (batch_size, seq_len, hidden_size) (batch_size, seq_len, 768)
        outputs_m = torch.stack(outputs_m)
        #no need to use query
        
        input_q=self.adapter(outputs_q)
        input_qa=self.alignment(outputs_qa)
        input_m=self.adapter1(outputs_m)
        # input_q=outputs_q
        # input_qa=outputs_qa

        
        next_item_embedding = self.user_llm(input_qa, target)
        
        # query_item_embedding=self.user_llm(input_q,target)
        """"
        将用户嵌入和下一个项目嵌入连接起来，然后通过一个线性层进行预测
        """
        # concat_q = torch.cat([next_item_embedding, outputs_q], dim=-1)
        concat_qa = torch.cat([next_item_embedding, input_q], dim=-1)
        output = self.prediction_head(concat_qa)
        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)
        
        
        # meomry loss
        concat_meo=torch.cat([next_item_embedding, input_m], dim=-1)
        # meo_output = self.meomry_head(concat_query)
        meo_output = self.meomry_head(concat_meo)
        # meo_output = self.meomry_head(query_item_embedding)
        meo_labels = meo_target.reshape(-1)
        
        meo_preds = (meo_output.reshape(-1))  # logit
        masked_meo_labels = meo_labels[mask].float()
        masked_meo_preds = meo_preds[mask]
        
        meo_output = self.loss(masked_meo_preds, masked_meo_labels)
        
        overall_loss = output.sum()+meo_output.sum()
        
        # overall_loss = output.sum()
        
        return overall_loss, self.m(preds), mask.sum()
        

