import torch
from tqdm import tqdm
import pickle
from transformers import BertTokenizer, BertModel, AutoModelForCausalLM, AutoTokenizer,AutoModel
import numpy as np
import math
import argparse
from torch.cuda.amp import autocast, GradScaler
import torch.nn as nn
from model.llmkt import load_llama_lora_model

def load_bert_model():
    model_name_or_path = "/home/shuaishuai/Gitroom/kg-bert/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)
    model = AutoModel.from_pretrained(model_name)
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    return tokenizer, model



def text_to_embedding(texts, model, tokenizer):
    outputs_embed = []
    for text in tqdm(texts, desc="Text to Embedding"):
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256)
        inputs = {key: val.to(model.device) for key, val in inputs.items()}
        with torch.no_grad():
            item_embeddings = model(**inputs).last_hidden_state.mean(dim=1)
        outputs_embed.append(item_embeddings.cpu())
    outputs = torch.cat(outputs_embed, dim=0)
    return outputs

def save_embeddings(texts, output_file):
    tokenizer, model = load_bert_model()
    model.to('cuda' if torch.cuda.is_available() else 'cpu')
    embeddings = text_to_embedding(texts, model, tokenizer)
    with open(output_file, 'wb') as f:
        pickle.dump(embeddings, f)
# Code reused from https://github.com/jennyzhang0215/DKVMN.git


def format_student_records(input_q, input_pid, target):
    formatted_records = []
    for i in range(len(input_q)):
        student_record = []
        for t in range(len(input_q[i])):
            qid = input_pid[i][t]
            cid = input_q[i][t]
            correct = target[i][t]
            status = "correctly" if correct == 1 else "incorrectly"
            if qid !=0:
                student_record.append(
                    f"The student has answered question QID={qid} involving concept CID={cid} {status} at timestep {t}.[ITEM]"
                )
            else:
                student_record.append(
                    "[None]"
                )
        formatted_records.append(student_record)
    return formatted_records

def format_student_query(input_q, input_pid):
    formatted_records = []
    for i in range(len(input_q)):
        student_record = []
        for t in range(len(input_q[i])):
            qid = input_pid[i][t]
            cid = input_q[i][t]
            if qid !=0:
                student_record.append(
                    f"Judge the student can answer question QID={qid} involving concept CID={cid} correctly or not.[ITEM]"
                )
            else:
                student_record.append(
                    "[None]"
                )
        formatted_records.append(student_record)
    return formatted_records

def format_student_records_concepts(input_q, target):
    formatted_records = []
    for i in range(len(input_q)):
        student_record = []
        for t in range(len(input_q[i])):
            cid = input_q[i][t]
            correct = target[i][t]
            status = "correctly" if correct == 1 else "incorrectly"
            if cid !=0:
                student_record.append(
                    f"The student has answered the question involving concept CID={cid} {status} at timestep {t}.[ITEM]"
                )
            else:
                student_record.append(
                    "[None]"
                )
        formatted_records.append(student_record)
    return formatted_records    
    
def format_student_query_concepts(input_q):
    formatted_records = []
    for i in range(len(input_q)):
        student_record = []
        for t in range(len(input_q[i])):
          
            cid = input_q[i][t]
            if cid !=0:
                student_record.append(
                    f"Judge the student can answer question involving concept CID={cid} correctly or not.[ITEM]"
                )
            else:
                student_record.append(
                    "[None]"
                )
        formatted_records.append(student_record)
    return formatted_records
 
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):

        """
        自动构造数据，直接传入文本列表
        """
        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 DATA(object):
    def __init__(self, n_question, seqlen, separate_char, name="data"):
        # In the ASSISTments2009 dataset:
        # param: n_queation = 110
        #        seqlen = 200
        self.separate_char = separate_char
        self.n_question = n_question
        self.seqlen = seqlen
    # data format
    # id, true_student_id
    # 1,1,1,1,7,7,9,10,10,10,10,11,11,45,54
    # 0,1,1,1,1,1,0,0,1,1,1,1,1,0,0

    def load_data(self, path):
        f_data = open(path, 'r')
        q_data = []
        qa_data = []
        idx_data = []
        for lineID, line in enumerate(f_data):
            line = line.strip()
            # lineID starts from 0
            if lineID % 3 == 0:
                student_id = lineID//3
            if lineID % 3 == 1:
                Q = line.split(self.separate_char)
                if len(Q[len(Q)-1]) == 0:
                    Q = Q[:-1]
                # print(len(Q))
            elif lineID % 3 == 2:
                A = line.split(self.separate_char)
                if len(A[len(A)-1]) == 0:
                    A = A[:-1]
                # print(len(A),A)

                # start split the data
                n_split = 1
                # print('len(Q):',len(Q))
                if len(Q) > self.seqlen:
                    n_split = math.floor(len(Q) / self.seqlen)
                    if len(Q) % self.seqlen:
                        n_split = n_split + 1
                # print('n_split:',n_split)
                for k in range(n_split):
                    question_sequence = []
                    answer_sequence = []
                    if k == n_split - 1:
                        endINdex = len(A)
                    else:
                        endINdex = (k+1) * self.seqlen
                    for i in range(k * self.seqlen, endINdex):
                        if len(Q[i]) > 0:
                            # int(A[i]) is in {0,1}
                            Xindex = int(Q[i]) + int(A[i]) * self.n_question
                            question_sequence.append(int(Q[i]))
                            answer_sequence.append(Xindex)
                        else:
                            print(Q[i])
                    #print('instance:-->', len(instance),instance)
                    q_data.append(question_sequence)
                    qa_data.append(answer_sequence)
                    idx_data.append(student_id)
        f_data.close()
        ### data: [[],[],[],...] <-- set_max_seqlen is used
        # convert data into ndarrays for better speed during training
        q_dataArray = np.zeros((len(q_data), self.seqlen))
        for j in range(len(q_data)):
            dat = q_data[j]
            q_dataArray[j, :len(dat)] = dat

        qa_dataArray = np.zeros((len(qa_data), self.seqlen))
        for j in range(len(qa_data)):
            dat = qa_data[j]
            qa_dataArray[j, :len(dat)] = dat
        # dataArray: [ array([[],[],..])] Shape: (3633, 200)
        return q_dataArray, qa_dataArray, np.asarray(idx_data)


class PID_DATA(object):
    def __init__(self, n_question,  seqlen, separate_char, name="data"):
        # In the ASSISTments2009 dataset:
        # param: n_queation = 110
        #        seqlen = 200
        self.separate_char = separate_char
        self.seqlen = seqlen
        self.n_question = n_question
    # data format
    # id, true_student_id
    # pid1, pid2, ...
    # 1,1,1,1,7,7,9,10,10,10,10,11,11,45,54
    # 0,1,1,1,1,1,0,0,1,1,1,1,1,0,0

    def load_data(self, path):
        f_data = open(path, 'r')
        q_data = []
        qa_data = []
        p_data = []
        user=[]
        order_id=0
        for lineID, line in enumerate(f_data):
            line = line.strip()
            # lineID starts from 0
            if lineID % 4 == 0:
                student_id = lineID//4
            if lineID % 4 == 2:
                Q = line.split(self.separate_char)
                if len(Q[len(Q)-1]) == 0:
                    Q = Q[:-1]
                # print(len(Q))
            if lineID % 4 == 1:
                P = line.split(self.separate_char)
                if len(P[len(P) - 1]) == 0:
                    P = P[:-1]

            elif lineID % 4 == 3:
                A = line.split(self.separate_char)
                if len(A[len(A)-1]) == 0:
                    A = A[:-1]
                # print(len(A),A)

                # start split the data
                n_split = 1
                # print('len(Q):',len(Q))
                if len(Q) > self.seqlen:
                    n_split = math.floor(len(Q) / self.seqlen)
                    if len(Q) % self.seqlen:
                        n_split = n_split + 1
                # print('n_split:',n_split)
                for k in range(n_split):
                    order_id=order_id+1
                    question_sequence = []
                    problem_sequence = []
                    answer_sequence = []
                    if k == n_split - 1:
                        endINdex = len(A)
                    else:
                        endINdex = (k+1) * self.seqlen
                    for i in range(k * self.seqlen, endINdex):
                        if len(Q[i]) > 0:
                            Xindex = int(Q[i]) + int(A[i]) * self.n_question
                            question_sequence.append(int(Q[i]))
                            problem_sequence.append(int(P[i]))
                            answer_sequence.append(Xindex)
                        else:
                            print(Q[i])
                    user.append(order_id)
                    q_data.append(question_sequence)
                    qa_data.append(answer_sequence)
                    p_data.append(problem_sequence)

        f_data.close()
        ### data: [[],[],[],...] <-- set_max_seqlen is used
        # convert data into ndarrays for better speed during training
        q_dataArray = np.zeros((len(q_data), self.seqlen))
        for j in range(len(q_data)):
            dat = q_data[j]
            q_dataArray[j, :len(dat)] = dat

        qa_dataArray = np.zeros((len(qa_data), self.seqlen))
        for j in range(len(qa_data)):
            dat = qa_data[j]
            qa_dataArray[j, :len(dat)] = dat

        p_dataArray = np.zeros((len(p_data), self.seqlen))
        for j in range(len(p_data)):
            dat = p_data[j]
            p_dataArray[j, :len(dat)] = dat
            
       
        user_dataArray = np.zeros((len(user), 1))
        for j in range(len(user)):
            dat = user[j]
            user_dataArray[j,0] = dat
            
        
        return user_dataArray, q_dataArray, qa_dataArray, p_dataArray


def text_to_embed(user_dataArray, q_dataArray, qa_dataArray, p_dataArray,status,length):
    
    tokenizer, model = load_bert_model()
    # tokenizer, model = load_llama_model()
    embed_lm=ItemLLM(model=model, tokenizer=tokenizer).to('cuda' if torch.cuda.is_available() else 'cpu')
    
    target = (qa_dataArray - 1) / params.n_question
    target_1 = np.floor(target)
 

    input_q = torch.from_numpy(q_dataArray).long()
    input_qa = torch.from_numpy(qa_dataArray).long()
    target = torch.from_numpy(target_1).float()
    input_pid = torch.from_numpy(p_dataArray).long()
    user_id = torch.from_numpy(user_dataArray).long()
    
    texts_qa=format_student_records(input_q, input_pid, target)
    texts_q=format_student_query(input_q, input_pid)
    
    cls_qa_embeddings = {}
    qa_order=0
    for texts in tqdm(texts_qa,desc="Text to Embedding"):
        qa_order=qa_order+1
        item_embeddings = embed_lm(texts)
        cls_qa_embeddings[qa_order]=item_embeddings
        
    dir1="data/embed/bert/assist09_"+str(status)+"_len_"+str(length)+"_qa_embeddings.pkl"
    with open(dir1, "wb") as f:
        pickle.dump(cls_qa_embeddings, f)
        
        
    cls_q_embeddings = {}
    q_order=0
    for texts in tqdm(texts_q,desc="Text to Embedding"):
        q_order=q_order+1
        item_embeddings = embed_lm(texts)
        cls_q_embeddings[q_order]=item_embeddings
    
    dir2="data/embed/bert/assist09_"+str(status)+"_len_"+str(length)+"_q_embeddings.pkl"
    with open(dir2, "wb") as f:
        pickle.dump(cls_q_embeddings, f)
    
def text_to_embed_concepts(user_dataArray, q_dataArray, qa_dataArray,status,dataser):
    
    tokenizer, model = load_bert_model()
    # tokenizer, model = load_llama_model()
    embed_lm=ItemLLM(model=model, tokenizer=tokenizer).to('cuda' if torch.cuda.is_available() else 'cpu')
    
    target = (qa_dataArray - 1) / params.n_question
    target_1 = np.floor(target)
 

    input_q = torch.from_numpy(q_dataArray).long()
    input_qa = torch.from_numpy(qa_dataArray).long()
    target = torch.from_numpy(target_1).float()
    # user_id = torch.from_numpy(user_dataArray).long()
    
    texts_qa=format_student_records_concepts(input_q, target)
    texts_q=format_student_query_concepts(input_q)
    
    cls_qa_embeddings = {}
    qa_order=0
    for texts in tqdm(texts_qa,desc="Text to Embedding"):
        qa_order=qa_order+1
        item_embeddings = embed_lm(texts)
        cls_qa_embeddings[qa_order]=item_embeddings
        
    dir1="data/embed/bert/"+str(dataset)+"_"+str(status)+"_"+str(status)+"_qa_embeddings.pkl"
    with open(dir1, "wb") as f:
        pickle.dump(cls_qa_embeddings, f)
        
        
    cls_q_embeddings = {}
    q_order=0
    for texts in tqdm(texts_q,desc="Text to Embedding"):
        q_order=q_order+1
        item_embeddings = embed_lm(texts)
        cls_q_embeddings[q_order]=item_embeddings
    
    dir2="data/embed/bert/"+str(dataset)+"_"+str(status)+"_q_embeddings.pkl"
    with open(dir2, "wb") as f:
        pickle.dump(cls_q_embeddings, f)

# if __name__=="__main__":
#     # tokenizer, model = load_bert_model()
#     tokenizer, model = load_llama_lora_model()
#     embed_lm=ItemLLM(model=model, tokenizer=tokenizer).to('cuda' if torch.cuda.is_available() else 'cpu')
#     texts=["[None]"]
#     item_embeddings = embed_lm(texts)
#     dir="data/embed/bert/None.pkl"
#     None_embeddings = {}
#     None_embeddings[10086]=item_embeddings
#     with open(dir, "wb") as f:
#         pickle.dump(None_embeddings, f)
    
#     with open("/root/autodl-tmp/LLM4KT/data/embed/bert/None.pkl", "rb") as f:
#         none_embeddings = pickle.load(f)
    
#     print(none_embeddings[10086].shape)

if __name__ == "__main__":
    # 示例文本数据
    # Parse Arguments
    parser = argparse.ArgumentParser()
    params = parser.parse_args()
    dataset = "assist2015"

    if dataset in {"assist2009_pid"}:
        params.n_question = 110
        params.batch_size = 1024
        params.seqlen = 50
        params.data_dir = 'data/'+dataset
        params.data_name = dataset
        params.n_pid = 16891

    if dataset in {"assist2017_pid"}:
        params.batch_size = 24
        params.seqlen = 200
        params.data_dir = 'data/'+dataset
        params.data_name = dataset
        params.n_question = 102
        params.n_pid = 3162

    if dataset in {"assist2015"}:
        params.n_question = 100
        params.batch_size = 24
        params.seqlen = 100
        params.n_pid   = 0
        params.data_dir = 'data/'+dataset
        params.data_name = dataset

    if dataset in {"statics"}:
        params.n_question = 1223
        params.batch_size = 24
        params.seqlen = 100
        params.n_pid = 0
        params.data_dir = 'data/'+dataset
        params.data_name = dataset
        
        
    if dataset in {"algebra05_pid"}:
        params.max_iter=2000
        params.n_question = 138
        params.n_pid = 1084
        params.batch_size = 24
        params.seqlen = 200
        params.data_dir = 'data/'+dataset
        params.data_name = dataset

    params.save = params.data_name
    params.load = params.data_name
    
    if params.n_pid > 0:    
        dat = PID_DATA(n_question=params.n_question,
                       seqlen=params.seqlen, separate_char=',')
        """
        训练数据加载
        """
        train_data_path = params.data_dir + "/" + \
            params.data_name + "_train1"+".csv"
        user_dataArray, q_dataArray, qa_dataArray, p_dataArray=dat.load_data(train_data_path)
        text_to_embed(user_dataArray, q_dataArray, qa_dataArray, p_dataArray,"train", params.seqlen)
        
        """
        验证数据加载
        """
        valid_data_path = params.data_dir + "/" + \
            params.data_name + "_valid1"+".csv"
        user_dataArray, q_dataArray, qa_dataArray, p_dataArray=dat.load_data(valid_data_path)
        text_to_embed(user_dataArray, q_dataArray, qa_dataArray, p_dataArray, "valid",params.seqlen)
        
        """
        测试数据加载
        """
        test_data_path = params.data_dir + "/" + \
            params.data_name + "_test1"+".csv"
        user_dataArray, q_dataArray, qa_dataArray, p_dataArray=dat.load_data(test_data_path)
        text_to_embed(user_dataArray, q_dataArray, qa_dataArray, p_dataArray, "test",params.seqlen)
        
        #llama
        # tokenizer, model = load_llama_lora_model()
        #bert
        tokenizer, model = load_bert_model()
        embed_lm=ItemLLM(model=model, tokenizer=tokenizer).to('cuda' if torch.cuda.is_available() else 'cpu')
        texts=["[None]"]
        item_embeddings = embed_lm(texts)
        dir="data/embed/bert/None.pkl"
        None_embeddings = {}
        None_embeddings[10086]=item_embeddings
        with open(dir, "wb") as f:
            pickle.dump(None_embeddings, f)
        
        # with open("/root/autodl-tmp/LLM4KT/data/embed/bert/None.pkl", "rb") as f:
        #     none_embeddings = pickle.load(f)
        
        # print(none_embeddings[10086].shape)
    else:
        dat = DATA(n_question=params.n_question,
                   seqlen=params.seqlen, separate_char=',')
        """
        训练数据加载
        """
        train_data_path = params.data_dir + "/" + \
            params.data_name + "_train1"+".csv"
        q_dataArray, qa_dataArray, user_dataArray=dat.load_data(train_data_path)
        text_to_embed_concepts(user_dataArray, q_dataArray, qa_dataArray,"train",dataset)
        
        """
        验证数据加载
        """
        valid_data_path = params.data_dir + "/" + \
            params.data_name + "_valid1"+".csv"
        q_dataArray, qa_dataArray, user_dataArray=dat.load_data(valid_data_path)
        text_to_embed_concepts(user_dataArray, q_dataArray, qa_dataArray, "valid",dataset)
        
        """
        测试数据加载
        """
        test_data_path = params.data_dir + "/" + \
            params.data_name + "_test1"+".csv"
        q_dataArray, qa_dataArray, user_dataArray=dat.load_data(test_data_path)
        text_to_embed_concepts(user_dataArray, q_dataArray, qa_dataArray, "test",dataset)
        
        #llama
        # tokenizer, model = load_llama_lora_model()
        #bert
        # tokenizer, model = load_bert_model()
        # embed_lm=ItemLLM(model=model, tokenizer=tokenizer).to('cuda' if torch.cuda.is_available() else 'cpu')
        # texts=["[None]"]
        # item_embeddings = embed_lm(texts)
        # dir="data/embed/bert/None.pkl"
        # None_embeddings = {}
        # None_embeddings[10086]=item_embeddings
        # with open(dir, "wb") as f:
        #     pickle.dump(None_embeddings, f)
        
        # with open("/root/autodl-tmp/LLM4KT/data/embed/bert/None.pkl", "rb") as f:
        #     none_embeddings = pickle.load(f)
        
        # print(none_embeddings[10086].shape)

    
    