from typing import List, Tuple, Any

import numpy as np
from transformers import AutoTokenizer
from bi_encoder import get_corpus, get_question

import torch

from model import CrossEncoderBert


device = "cuda" if torch.cuda.is_available() else "cpu"

model = CrossEncoderBert()
model.model.resize_token_embeddings(len(model.tokenizer))
model.load_state_dict(torch.load('model/torch_model', map_location=torch.device(device)))
model.tokenizer = AutoTokenizer.from_pretrained('model/tokenizer')
model.to(device)


def get_range_answers(
                      context: str,
                      question: str,
                      num_answers: int = 5) -> list[str]:

    corpus = get_corpus(context, question)
    context_question = f'{context} [Cont_token] {question}'
    context_questions = [context_question] * len(corpus)
    tokenized_texts = model.tokenizer(
        context_questions,
        corpus,
        max_length=512,
        padding=True,
        truncation=True,
        return_tensors='pt'
    ).to(device)

    with torch.no_grad():
        ce_scores = model(tokenized_texts['input_ids'],
                          tokenized_texts['attention_mask']).squeeze(-1)
        ce_scores = torch.sigmoid(ce_scores)

    scores = ce_scores.cpu().numpy()
    scores_ix = np.argsort(scores)[::-1]
    best_answers = []
    for idx in scores_ix[:num_answers]:
        best_answers.append((scores[idx], corpus[idx]))

    best_answers = [str(x[1]) for x in best_answers]
    return best_answers


def get_best_answer(
        context: str,
        question: str
) -> str:
    return get_range_answers(context, question, 1)