from huggingface_hub import HfApi, ModelFilter
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
from transformers.tokenization_utils_base import BatchEncoding
from transformers.modeling_outputs import MaskedLMOutput

# Function to fetch suitable ESM models from HuggingFace Hub
def get_models() -> list[None|str]:
    """Fetch suitable ESM models from HuggingFace Hub."""
    if not any(
        out := [
            m.modelId for m in HfApi().list_models(
                filter=ModelFilter(
                    author="facebook", model_name="esm", task="fill-mask"
                ), 
                sort="lastModified", 
                direction=-1
            )
        ]
    ):
        raise RuntimeError("Error while retrieving models from HuggingFace Hub")
    return out

# Class to wrap ESM models
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()
            # Check if CUDA is available and if so, use it
            if torch.cuda.is_available():
                self.model = self.model.cuda()

    def tokenise(self, input: str) -> BatchEncoding:
        """Convert input string to batch of tokens."""
        return self.batch_converter(input, return_tensors="pt")

    def __call__(self, batch_tokens: torch.Tensor, **kwargs) -> MaskedLMOutput:
        """Run model on batch of tokens."""
        return self.model(batch_tokens, **kwargs)

    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."""
            # Extract wild type, index and mutant type from the row
            wt, idx, mt = row[0], int(row[1:-1])-1, row[-1]
            # Calculate the score as the difference between the token probabilities of the mutant type and the wild type
            score = token_probs[0, 1+idx, self[mt]] - token_probs[0, 1+idx, self[wt]]
            return score.item()

        # Tokenise the sequence data
        batch_tokens = self.tokenise(data.seq).input_ids

        # Calculate the token probabilities without updating the model parameters
        with torch.no_grad():
            token_probs = torch.log_softmax(self(batch_tokens).logits, dim=-1)
        # Store the token probabilities in the data
        data.token_probs = token_probs.cpu().numpy()

        # If the scoring strategy starts with "masked-marginals"
        if data.scoring_strategy.startswith("masked-marginals"):
            all_token_probs = []
            # For each token in the batch
            for i in range(batch_tokens.size()[1]):
                # If the token is in the list of residues
                if i in data.resi:
                    # Clone the batch tokens and mask the current token
                    batch_tokens_masked = batch_tokens.clone()
                    batch_tokens_masked[0, i] = self['<mask>']
                    # Calculate the masked token probabilities
                    with torch.no_grad():
                        masked_token_probs = torch.log_softmax(
                            self(batch_tokens_masked).logits, dim=-1
                        )
                else:
                    # If the token is not in the list of residues, use the original token probabilities
                    masked_token_probs = token_probs
                # Append the token probabilities to the list
                all_token_probs.append(masked_token_probs[:, i])
            # Concatenate all token probabilities
            token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)

        # Apply the label_row function to each row of the substitutions dataframe
        data.out[self.model_name] = data.sub.apply(
            lambda row: label_row(
                row['0'],
                token_probs,
            ),
            axis=1,
        )