from agentdriver.llm_core.timeout import timeout
from transformers import (BertModel, BertTokenizer, AutoTokenizer, 
                          DPRContextEncoder, AutoModel, RealmEmbedder, RealmForOpenQA)
from agentdriver.memory.common_sense_memory import CommonSenseMemory
from agentdriver.memory.experience_memory import ExperienceMemory
# from embedder.train_classification_retriever import ClassificationNetwork
# from embedder.train_contrastive_retriever import TripletNetwork
import torch
from torch import nn

class TripletNetwork(nn.Module):
    def __init__(self):
        super(TripletNetwork, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output  # Using the pooled output of BERT
        return pooled_output

class ClassificationNetwork(nn.Module):
    def __init__(self, num_labels):
        super(ClassificationNetwork, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output

class DPRNetwork(nn.Module):
    def __init__(self):
        super(DPRNetwork, self).__init__()
        self.bert = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base").to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output

class ANCENetwork(nn.Module):
    def __init__(self):
        super(ANCENetwork, self).__init__()
        self.bert = AutoModel.from_pretrained("castorini/ance-dpr-question-multi").to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output
    
class REALMNetwork(nn.Module):
    def __init__(self):
        super(REALMNetwork, self).__init__()
        self.bert = RealmEmbedder.from_pretrained("google/realm-cc-news-pretrained-embedder").realm.to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output

class BGENetwork(nn.Module):
    def __init__(self):
        super(BGENetwork, self).__init__()
        self.bert = AutoModel.from_pretrained("BAAI/bge-large-en").to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output

class ORQANetwork(nn.Module):
    def __init__(self):
        super(ORQANetwork, self).__init__()
        self.bert = RealmForOpenQA.from_pretrained("google/realm-orqa-nq-openqa").embedder.realm.to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output

class BM25Network(nn.Module):
    def __init__(self):
        super(BM25Network, self).__init__()
        self.bert = AutoModel.from_pretrained("facebook/spar-wiki-bm25-lexmodel-context-encoder").to("cuda")

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        # pooled_output = self.dropout(pooled_output)
        # logits = self.classifier(pooled_output)
        return pooled_output


class MemoryAgent:
    def __init__(self, data_path, model_name="gpt-3.5-turbo-0613", verbose=False, compare_perception=False, embedding="Linear", args=None) -> None:
        self.model_name = model_name
        self.common_sense_memory = CommonSenseMemory()
        self.embedding = embedding
        if self.embedding == "Contrastive":
            embedder_dir = 'RAG/embedder/contrastive_embedder_user_random_diverse/checkpoint-300'
            self.embedding_model = TripletNetwork().to("cuda")
            self.embedding_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
            # Load the weights
            self.embedding_model.load_state_dict(torch.load(embedder_dir + "/pytorch_model.bin"))
            self.embedding_model.eval()  # for inference

        elif self.embedding == "Classification":
            embedder_dir = 'RAG/embedder/classification_embedder_user/checkpoint-500'
            num_labels = 11
            self.embedding_model = ClassificationNetwork(num_labels=num_labels).to("cuda")
            self.embedding_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
            # Load the weights
            self.embedding_model.load_state_dict(torch.load(embedder_dir + "/pytorch_model.bin"))
            self.embedding_model.eval()  # for inference
        
        elif self.embedding == "facebook/dpr-ctx_encoder-single-nq-base":

            # load retriever
            self.embedding_model = DPRNetwork().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
            # Load the weights
            self.embedding_model.eval()  # for inference

        elif self.embedding == "castorini/ance-dpr-question-multi":

            # load retriever
            self.embedding_model = ANCENetwork().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("castorini/ance-dpr-question-multi")
            # Load the weights
            self.embedding_model.eval()  # for inference

        elif self.embedding == "bge-large-en":

            # load retriever
            self.embedding_model = BGENetwork().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-en")
            # Load the weights
            self.embedding_model.eval()

        elif self.embedding == "realm-cc-news-pretrained-embedder":
                
            # load retriever
            self.embedding_model = REALMNetwork().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("google/realm-cc-news-pretrained-embedder")
            # Load the weights
            self.embedding_model.eval()

        elif self.embedding == "realm-orqa-nq-openqa":
                    
            # load retriever
            self.embedding_model = ORQANetwork().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("google/realm-orqa-nq-openqa")
            # Load the weights
            self.embedding_model.eval()

        elif self.embedding == "spar-wiki-bm25-lexmodel-context-encoder":

            # load retriever
            self.embedding_model = BM25Network().to("cuda")
            self.embedding_tokenizer = AutoTokenizer.from_pretrained("facebook/spar-wiki-bm25-lexmodel-context-encoder")
            # Load the weights
            self.embedding_model.eval()
        
        else:
            self.embedding_model = None
            self.embedding_tokenizer = None

        self.experience_memory = ExperienceMemory(data_path, model_name=self.model_name, verbose=verbose, compare_perception=compare_perception, embedding=self.embedding, embedding_model=self.embedding_model, embedding_tokenizer=self.embedding_tokenizer, args=args)
        self.verbose = verbose


    def retrieve(self, working_memory):
        raise NotImplementedError
    
    def retrieve_common_sense_memory(self, knowledge_types: list = None):
        return self.common_sense_memory.retrieve(knowledge_types=knowledge_types)

    def retrieve_experience_memory(self, working_memory, embedding):
        return self.experience_memory.retrieve(working_memory)

    def insert(self, working_memory):
        raise NotImplementedError

    def update(self, working_memory):
        raise NotImplementedError

    @timeout(15)
    def run(self, working_memory):
        common_sense_prompts = self.retrieve_common_sense_memory()
        experience_prompt = self.retrieve_experience_memory(working_memory, self.embedding)

        return common_sense_prompts, experience_prompt