from huggingface_hub import HfApi, ModelFilter
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

    # fetch suitable ESM models from HuggingFace Hub
MODELS = [m.modelId for m in HfApi().list_models(filter=ModelFilter(author="facebook", model_name="esm", task="fill-mask"), sort="lastModified", direction=-1)]
if not any(MODELS):
    raise RuntimeError("Error while retrieving models from HuggingFace Hub")

class Model:
    """Wrapper for ESM models"""
    def __init__(self, model_name:str=""):
        "load selected model and tokenizer"
        self.model_name = model_name
        if model_name:
            self.model = AutoModelForMaskedLM.from_pretrained(model_name)
            self.batch_converter = AutoTokenizer.from_pretrained(model_name)
            self.alphabet = self.batch_converter.get_vocab()
            if torch.cuda.is_available():
                self.model = self.model.cuda()

    def __rshift__(self, batch_tokens:torch.Tensor) -> torch.Tensor:
        "run model on batch of tokens"
        return self.model(batch_tokens)["logits"]
    
    def __lshift__(self, input:str) -> torch.Tensor:
        "convert input string to batch of tokens"
        return self.batch_converter(input, return_tensors="pt")["input_ids"]

    def __getitem__(self, key:str) -> int:
        "get token ID from character"
        return self.alphabet[key]
    
    def run_model(self, data):
        "run model on data"
        def label_row(row, token_probs):
            "label row with score"
            wt, idx, mt = row[0], int(row[1:-1])-1, row[-1]
            score = token_probs[0, 1+idx, self[mt]] - token_probs[0, 1+idx, self[wt]]
            return score.item()
        
        batch_tokens = self<<data.seq

            # run model with selected scoring strategy (info thereof available in the original ESM paper)
        if data.scoring_strategy.startswith("wt-marginals"):
            with torch.no_grad():
                token_probs = torch.log_softmax(self>>batch_tokens, dim=-1)
            data.out[self.model_name] = data.sub.apply(
                lambda row: label_row(
                    row['0'],
                    token_probs,
                ),
                axis=1,
            )
        elif data.scoring_strategy.startswith("masked-marginals"):
            all_token_probs = []
            for i in range(batch_tokens.size()[1]):
                batch_tokens_masked = batch_tokens.clone()
                batch_tokens_masked[0, i] = self['<mask>']
                with torch.no_grad():
                    token_probs = torch.log_softmax(
                        self>>batch_tokens_masked, dim=-1
                    )
                all_token_probs.append(token_probs[:, i])
            token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
            data.out[self.model_name] = data.sub.apply(
                lambda row: label_row(
                    row['0'],
                    token_probs,
                ),
                axis=1,
            )