import pandas as pd
import pickle
import torch
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
bert_model = AutoModel.from_pretrained("distilbert-base-uncased")


def mean_pool(token_embeds: torch.tensor, attention_mask: torch.tensor) -> torch.tensor:
    in_mask = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
    pool = torch.sum(token_embeds * in_mask, 1) / torch.clamp(in_mask.sum(1), min=1e-9)
    return pool


def encode(input_texts: list[str], tokenizer: AutoTokenizer, model: AutoModel, device: str = "cpu"
) -> torch.tensor:

    model.eval()
    tokenized_texts = tokenizer(input_texts, max_length=512,
                                padding='max_length', truncation=True, return_tensors="pt")
    token_embeds = model(tokenized_texts["input_ids"].to(device),
                         tokenized_texts["attention_mask"].to(device)).last_hidden_state
    pooled_embeds = mean_pool(token_embeds, tokenized_texts["attention_mask"].to(device))
    return pooled_embeds


with open('data/sentences.pkl', 'rb') as f:
    sentences = pickle.load(f)

with open('data/corpus.pkl', 'rb') as f:
    corpus = pickle.load(f)

df = pd.DataFrame.from_dict(sentences)
df['corpus'] = corpus


def get_question(context: str, question: str):
    cont_quest = f"{context} [Cont_token] {question}"
    pooled_embeds = encode(cont_quest, tokenizer, bert_model, "cpu")
    pooled_embeds = pooled_embeds.cpu().detach().numpy()
    return pooled_embeds


def cosine_sim(question, embed):
    return cosine_similarity(question, embed)[0][0]


def get_corpus(context: str, question: str):
    question_embed = get_question(context, question)
    df['cosine_similarity'] = df.apply(lambda x: cosine_sim(question_embed, x['embeds']), axis=1)
    corp = df.sort_values(by=['cosine_similarity'], ascending=False).head(10)['corpus'].tolist()
    return corp