import os

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

cross_encoder = None
cross_enc_tokenizer = None

TOP_K_RERANK = os.getenv("TOP_K_RERANK", 40)


@torch.no_grad()
def rerank_with_cross_encoder(cross_enc_name, documents, query):
    if cross_enc_name is None or len(documents) <= 1:
        return documents

    global cross_encoder, cross_enc_tokenizer
    if cross_encoder is None or cross_encoder.name_or_path != cross_enc_name:
        cross_encoder = AutoModelForSequenceClassification.from_pretrained(cross_enc_name)
        cross_encoder.eval()
        cross_enc_tokenizer = AutoTokenizer.from_pretrained(cross_enc_name)

    features = cross_enc_tokenizer(
        [query] * len(documents), documents, padding=True, truncation=True, return_tensors="pt"
    )
    scores = cross_encoder(**features).logits.squeeze()
    ranks = torch.argsort(scores, descending=True)
    documents = [documents[i] for i in ranks[:TOP_K_RERANK]]
    return documents