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 ItemLLM_Fast(nn.Module):
    def __init__(self, model, tokenizer):
        super(ItemLLM_Fast, self).__init__()
        self.llm = model
        self.tokenizer = tokenizer
        
        # 1. 合并所有特殊标记
        self.special_tokens = {
            'item': "[ITEM]",
            'none': "[None]",
            'query': "[QUERY]",
            'pad': "[None1]",
            'meo': "[MEO]",
            'pad2': "[None2]"
        }
        # self.special_tokens = {
        #     'item': "[ITEM]",
        #     'none': "[None]",
        #     'query': "[QUERY]",
        #     'pad': "[None1]",
        # }
        
        # 2. 一次性添加所有特殊标记
        self.tokenizer.add_special_tokens({"additional_special_tokens": list(self.special_tokens.values())})
        self.llm.resize_token_embeddings(len(self.tokenizer))
        
        # 3. 预计算所有特殊标记的ID
        self.token_ids = {
            k: self.tokenizer.convert_tokens_to_ids(v) 
            for k, v in self.special_tokens.items()
        }
        
        # 4. 添加缓存机制
        self.cache = {}
        self.cache_size = 1000  # 可调整的缓存大小
        
 
        
        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)
        
        # 9. 使用attention mask优化计算
        attention_mask = inputs["attention_mask"]
        
        # 10. 一次性前向传播
        # with torch.cuda.amp.autocast():  # 使用混合精度
        outputs = self.llm(**inputs)
        last_hidden_state = outputs.last_hidden_state
        
        # 计算每句话中特殊token与其他token的余弦相似度
        batch_size = last_hidden_state.size(0)
        seq_length = last_hidden_state.size(1)
        
        # 存储所有句子的余弦相似度结果
        all_similarities = []
        
        for batch_idx in range(batch_size):
            sentence_similarities = {}
            
            # 获取当前句子的所有token嵌入
            sentence_embeddings = last_hidden_state[batch_idx]  # [seq_length, hidden_dim]
            
            # 获取当前句子的所有token字符
            sentence_token_list = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][batch_idx])
            
            # 计算[ITEM] token与之前token的余弦相似度
            item_positions = (inputs["input_ids"][batch_idx] == self.token_ids['item']).nonzero(as_tuple=True)[0]
            if len(item_positions) > 0:
                item_similarities = []
                item_tokens = []
                item_preceding_tokens = []
                for item_pos in item_positions:
                    if item_pos > 0:  # 确保有前面的token
                        item_embedding = sentence_embeddings[item_pos]  # [hidden_dim]
                        # 只计算与之前位置的余弦相似度
                        preceding_embeddings = sentence_embeddings[:item_pos]  # [item_pos, hidden_dim]
                        similarities = F.cosine_similarity(
                            item_embedding.unsqueeze(0),  # [1, hidden_dim]
                            preceding_embeddings,  # [item_pos, hidden_dim]
                            dim=1
                        )  # [item_pos]
                        item_similarities.append(similarities.tolist())
                        # 保存对应位置的token字符
                        token_id = inputs["input_ids"][batch_idx][item_pos].item()
                        token_char = self.tokenizer.decode([token_id])
                        item_tokens.append(token_char)
                        # 保存之前位置对应的token字符
                        item_preceding_tokens.append(sentence_token_list[:item_pos])
                if item_similarities:  # 只有当有相似度数据时才添加
                    sentence_similarities['item'] = {
                        'similarities': item_similarities,
                        'tokens': item_tokens,
                        'preceding_tokens': item_preceding_tokens
                    }
            
            # 计算[QUERY] token与之前token的余弦相似度
            query_positions = (inputs["input_ids"][batch_idx] == self.token_ids['query']).nonzero(as_tuple=True)[0]
            if len(query_positions) > 0:
                query_similarities = []
                query_tokens = []
                query_preceding_tokens = []
                for query_pos in query_positions:
                    if query_pos > 0:  # 确保有前面的token
                        query_embedding = sentence_embeddings[query_pos]  # [hidden_dim]
                        # 只计算与之前位置的余弦相似度
                        preceding_embeddings = sentence_embeddings[:query_pos]  # [query_pos, hidden_dim]
                        similarities = F.cosine_similarity(
                            query_embedding.unsqueeze(0),  # [1, hidden_dim]
                            preceding_embeddings,  # [query_pos, hidden_dim]
                            dim=1
                        )  # [query_pos]
                        query_similarities.append(similarities.tolist())
                        # 保存对应位置的token字符
                        token_id = inputs["input_ids"][batch_idx][query_pos].item()
                        token_char = self.tokenizer.decode([token_id])
                        query_tokens.append(token_char)
                        # 保存之前位置对应的token字符
                        query_preceding_tokens.append(sentence_token_list[:query_pos])
                if query_similarities:  # 只有当有相似度数据时才添加
                    sentence_similarities['query'] = {
                        'similarities': query_similarities,
                        'tokens': query_tokens,
                        'preceding_tokens': query_preceding_tokens
                    }
            
            # 计算[MEO] token与之前token的余弦相似度
            meo_positions = (inputs["input_ids"][batch_idx] == self.token_ids['meo']).nonzero(as_tuple=True)[0]
            if len(meo_positions) > 0:
                meo_similarities = []
                meo_tokens = []
                meo_preceding_tokens = []
                for meo_pos in meo_positions:
                    if meo_pos > 0:  # 确保有前面的token
                        meo_embedding = sentence_embeddings[meo_pos]  # [hidden_dim]
                        # 只计算与之前位置的余弦相似度
                        preceding_embeddings = sentence_embeddings[:meo_pos]  # [meo_pos, hidden_dim]
                        similarities = F.cosine_similarity(
                            meo_embedding.unsqueeze(0),  # [1, hidden_dim]
                            preceding_embeddings,  # [meo_pos, hidden_dim]
                            dim=1
                        )  # [meo_pos]
                        meo_similarities.append(similarities.tolist())
                        # 保存对应位置的token字符
                        token_id = inputs["input_ids"][batch_idx][meo_pos].item()
                        token_char = self.tokenizer.decode([token_id])
                        meo_tokens.append(token_char)
                        # 保存之前位置对应的token字符
                        meo_preceding_tokens.append(sentence_token_list[:meo_pos])
                if meo_similarities:  # 只有当有相似度数据时才添加
                    sentence_similarities['meo'] = {
                        'similarities': meo_similarities,
                        'tokens': meo_tokens,
                        'preceding_tokens': meo_preceding_tokens
                    }
            
            # 只有当有相似度数据时才添加到结果中
            if sentence_similarities:
                all_similarities.append(sentence_similarities)
        
        # 计算信息熵
        import numpy as np
        
        # 计算所有特殊token的信息熵总和
        total_entropy = 0.0
        entropy_tensors = []  # 存储每个相似度列表的熵张量
        
        for batch_idx in range(batch_size):
            # 获取当前batch的相似度数据
            if batch_idx < len(all_similarities):
                batch_data = all_similarities[batch_idx]
                
                # 计算ITEM token的信息熵
                if 'item' in batch_data and batch_data['item']['similarities']:
                    for similarities in batch_data['item']['similarities']:
                        similarities_array = torch.tensor(similarities, device=last_hidden_state.device, dtype=torch.float32)
                        # similarities_array = torch.abs(similarities_array)
                        similarities_array = (similarities_array + 1) / 2 
                        if torch.sum(similarities_array) > 0:
                            probabilities = similarities_array / torch.sum(similarities_array)
                            # 计算信息熵: H = -sum(p * log2(p))
                            entropy = -torch.sum(probabilities * torch.log2(probabilities + 1e-10))
                            total_entropy += entropy.item()
                            entropy_tensors.append(entropy)
                
                # 计算QUERY token的信息熵
                if 'query' in batch_data and batch_data['query']['similarities']:
                    for similarities in batch_data['query']['similarities']:
                        similarities_array = torch.tensor(similarities, device=last_hidden_state.device, dtype=torch.float32)
                        # similarities_array = torch.abs(similarities_array)
                        similarities_array = (similarities_array + 1) / 2 
                        if torch.sum(similarities_array) > 0:
                            probabilities = similarities_array / torch.sum(similarities_array)
                            entropy = -torch.sum(probabilities * torch.log2(probabilities + 1e-10))
                            total_entropy += entropy.item()
                            entropy_tensors.append(entropy)
                
                # 计算MEO token的信息熵
                if 'meo' in batch_data and batch_data['meo']['similarities']:
                    for similarities in batch_data['meo']['similarities']:
                        similarities_array = torch.tensor(similarities, device=last_hidden_state.device, dtype=torch.float32)
                        # similarities_array = torch.abs(similarities_array)
                        similarities_array = (similarities_array + 1) / 2 
                        if torch.sum(similarities_array) > 0:
                            probabilities = similarities_array / torch.sum(similarities_array)
                            entropy = -torch.sum(probabilities * torch.log2(probabilities + 1e-10))
                            total_entropy += entropy.item()
                            entropy_tensors.append(entropy)
        
        # 将信息熵转换为可用于反向传播的张量
        if entropy_tensors:
            # 计算平均信息熵
            avg_entropy = torch.stack(entropy_tensors).mean()
            # 将信息熵作为模型的属性，可以在训练时使用
            self.current_entropy = avg_entropy
        else:
            self.current_entropy = torch.tensor(0.0, device=last_hidden_state.device)
        

        # 11. 并行提取所有特殊标记的嵌入
        embeddings = {}
        for token_name, token_id in self.token_ids.items():
            token_index = (inputs["input_ids"] == token_id).nonzero(as_tuple=True)
            token_embeddings = last_hidden_state[token_index]
            embeddings[token_name] = token_embeddings
        
        # 12. 合并相关的嵌入
        item_embeddings = torch.cat((embeddings['item'], embeddings['none']), 0)
        query_embeddings = torch.cat((embeddings['query'], embeddings['none']), 0)
        meo_embeddings = torch.cat((embeddings['meo'], embeddings['none']), 0)
   
        # 返回嵌入（不返回相似度）
        return item_embeddings, query_embeddings, meo_embeddings
        # return item_embeddings, query_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



# 该模型使用了 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.item_llm = ItemLLM_Fast(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 for g1,g2 in zip(input_q, input_qa)]
        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)
            # item_embeddings,query_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
        # outputs_qa,outputs_q,outputs_m = self.item_llm(merged_texts)
        
        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()+0.5*meo_output.sum()
        
     
        entropy_loss = self.item_llm.current_entropy
        # 使用较小的权重系数，避免信息熵损失主导整个训练过程
        entropy_weight = 0.1
        overall_loss = overall_loss + entropy_weight * entropy_loss
 
        # overall_loss = output.sum()
        
        return overall_loss, self.m(preds), mask.sum(), self.m(meo_preds)
        
        

