"""Module for calculating masked position probabilities using ESM models and Modal.

This module uses functions and classes to compute the probabilities at masked
positions in sequences using ESM models.
"""

from __future__ import annotations

import torch

from apb.biomodal.esm2.constants import console


def _get_model_device(model: torch.nn.Module):
    return next(model.parameters()).device


def calculate_pseudo_log_likelihood(
    token_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    token_mask: torch.Tensor,
    model: torch.nn.Module,
    mask_token_id: int,
    quiet: bool = True,
) -> torch.Tensor:
    device = _get_model_device(model)
    token_ids = token_ids.to(device)
    attention_mask = attention_mask.to(device)
    token_mask = token_mask.to(device)
    num_seqs, input_length = token_ids.size()
    pseudo_log_likelihood = torch.zeros(num_seqs, device=device)

    for i in range(input_length):
        if not quiet:
            console.print(f"Processed {i}/{input_length} token positions.")

        if token_mask[:, i].sum() == 0:
            continue

        masked_token_ids = token_ids.clone()
        masked_token_ids[:, i] = mask_token_id

        with torch.no_grad():
            output = model(masked_token_ids, attention_mask=attention_mask)
            logits_i = output.logits[:, i, :]
            logit_prob = torch.nn.functional.log_softmax(logits_i, dim=-1)

            position_log_probs = logit_prob[
                torch.arange(0, num_seqs, device=device), token_ids[:, i]
            ]
            pseudo_log_likelihood += position_log_probs * token_mask[:, i]

            del output, logits_i, logit_prob, position_log_probs
            torch.cuda.empty_cache()

    return pseudo_log_likelihood.cpu()


def calculate_pseudo_perplexity(plls: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
    return torch.exp(-plls / lengths)
