"""Module for calculating protein sequence perplexity using ESM models and Modal.

Possible models:
    - facebook/esm2_t6_8M_UR50D
    - facebook/esm2_t12_35M_UR50D
    - facebook/esm2_t30_150M_UR50D
    - facebook/esm2_t33_650M_UR50D
    - facebook/esm2_t36_3B_UR50D
    - facebook/esm2_t48_15B_UR50D
"""

from __future__ import annotations

import pandas as pd
import torch

from apb.biomodal.esm2.batching import create_length_batched_sequences
from apb.biomodal.esm2.cache_model import download_and_cache_model
from apb.biomodal.esm2.constants import (
    HOURS,
    MAX_SEQ_LENGTH,
    console,
)
from apb.biomodal.esm2.model import EsmModel, app
from apb.biomodal.esm2.ops import calculate_pseudo_perplexity


@app.function(
    gpu=None,
    memory=1024,
    cpu=1.0,
    timeout=24 * HOURS,
    scaledown_window=60,
)
def calculate_esm_pppl(
    sequences: dict[str, str],
    model_name: str,
    gpu: str,
    num_gpus: int,
    batch_size: int,
) -> torch.Tensor:
    """Calculates pseudo-perplexity for a dictionary of protein sequences.

    Args:
        sequences:
            Dictionary of protein sequences.
        model_name:
            The name of the ESM model to use for perplexity calculation.
        gpu:
            The type of GPU to use. Can be "any", "A100", etc.
        num_gpus:
            The number of GPU workers to spawn.
        batch_size:
            The number of sequences to process in each batch.

    Returns:
        A tensor of pseudo-perplexity values. Order matches input dictionary.
    """
    download_and_cache_model.remote(model_name, force_download=False)

    console.print(f"Using model {model_name}.")

    EsmModelWithGpu = EsmModel.with_options(  # type: ignore
        gpu=gpu, max_containers=num_gpus
    )
    worker = EsmModelWithGpu(model_name)

    sequence_batches, index_batches = create_length_batched_sequences(sequences, batch_size)

    console.print(f"Submitting a total of {len(sequence_batches)} batches.")

    _results = worker.pseudo_log_likelihood.map(sequence_batches, order_outputs=True)

    plls_sorted = torch.cat(list(_results), dim=0)

    sorted_sequences: list[str] = []
    for sequence_batch in sequence_batches:
        sorted_sequences.extend(sequence_batch)

    sorted_indices: list[int] = []
    for index_batch in index_batches:
        sorted_indices.extend(index_batch)

    pppls_sorted = calculate_pseudo_perplexity(
        plls_sorted, torch.tensor([min(len(seq), MAX_SEQ_LENGTH) for seq in sorted_sequences])
    )

    pppls = torch.zeros_like(pppls_sorted)
    for i in range(len(pppls_sorted)):
        pppls[sorted_indices[i]] = pppls_sorted[i]

    return pppls


def _get_pppl_frame(sequences: dict[str, str], pppl_tensor: torch.Tensor) -> pd.DataFrame:
    pppl_dict = {seq_id: pppl.item() for seq_id, pppl in zip(sequences, pppl_tensor, strict=True)}
    return (
        pd.Series(pppl_dict, name="pppl")
        .to_frame()
        .reset_index()
        .rename(columns={"index": "seq_id"})
    )


@app.function(
    gpu=None,
    memory=1024,
    cpu=1.0,
    timeout=24 * HOURS,
)
def pseudo_perplexity(
    seqs: dict[str, str],
    model_name: str,
    gpu: str,
    num_gpus: int,
    batch_size: int,
) -> pd.DataFrame:
    """Remote-callable function for protein sequence perplexity analysis.

    This function can be invoked remotely (after deployment) using:

    ```python
    import modal

    from biotite.sequence.io.fasta import FastaFile
    fasta_file = "example.fa"
    seqs = {defline: seq for defline, seq in FastaFile.read_iter(fasta_file)}

    pppl = modal.Function.from_name("esm2", "pseudo_perplexity").remote(
        seqs=seqs,
        model_name="facebook/esm2_t6_8M_UR50D",
        gpu="T4",
        batch_size=8,
        num_gpus=1,
    )

    print(pppl.head())
    ```
    """
    if any(len(x) > MAX_SEQ_LENGTH for x in seqs.values()):
        console.print(
            f"Some sequences passed are greater than {MAX_SEQ_LENGTH}, the max sequence length. "
            f"These sequences will be truncated, meaning PPPL will only be calculated for portions "
            f"of these sequences."
        )

    pppl_tensor = calculate_esm_pppl.local(seqs, model_name, gpu, num_gpus, batch_size)

    return _get_pppl_frame(seqs, pppl_tensor)
