# Description: Using Facebook's Faiss library to perform semantic search according to the query
# Reference: https://deepnote.com/blog/semantic-search-using-faiss-and-mpnet

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
import faiss

class SemanticEmbedding:

    def __init__(self, model_name="sentence-transformers/all-mpnet-base-v2"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)

    # Mean Pooling - Take attention mask into account for correct averaging
    def mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[
            0
        ]  # First element of model_output contains all token embeddings
        input_mask_expanded = (
            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        )
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
            input_mask_expanded.sum(1), min=1e-9
        )

    def get_embedding(self, sentences):
        # Tokenize sentences
        encoded_input = self.tokenizer(
            sentences, padding=True, truncation=True, return_tensors="pt"
        )
        with torch.no_grad():
            model_output = self.model(**encoded_input)
        # Perform pooling
        sentence_embeddings = self.mean_pooling(
            model_output, encoded_input["attention_mask"]
        )

        # Normalize embeddings
        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
        return sentence_embeddings.detach().numpy()


class FaissForQuerySearch:

    def __init__(self, model, dim=768):
        self.index = faiss.IndexFlatIP(dim)
        # Maintaining the document data
        self.doc_map = dict()
        self.model = model
        self.ctr = 0
        self.uuid = []
        self.labels = []

    def search_query(self, query, k=1):
        D, I = self.index.search(self.model.get_embedding(query), k)
        return [
            self.labels[idx] for idx, score in zip(I[0], D[0]) if idx in self.doc_map
        ]

    def add_summary(self, document_text, id, predicted_label):
        self.index.add((self.model.get_embedding(document_text)))  # index
        self.uuid.append(id)  # appending the uuid
        self.labels.append(predicted_label)  # appending the predicted label
        self.doc_map[self.ctr] = document_text  # store the original document text
        self.ctr += 1