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

from __future__ import annotations
from collections.abc import Generator

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, console
from apb.biomodal.esm2.model import EsmModel, app


def calculate_embeddings(
    sequences: dict[str, str],
    model_name: str,
    gpu: str,
    num_gpus: int,
    batch_size: int,
    pooling: str = "mean",
) -> Generator[tuple[torch.Tensor, list[str], list[int]], None, None]:
    """Calculates embeddings for a dictionary of protein sequences.

    Args:
        sequences:
            Dictionary mapping sequence IDs to protein sequences.
        model_name:
            The name of the ESM model to use.
        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.
        pooling:
            Pooling strategy. "mean" for mean-pooled embeddings, "none" for per-residue.

    Yields:
        Tuples of (embeddings_tensor, seq_ids, seq_lengths) for each batch.
    """
    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)

    seq_ids = list(sequences.keys())
    seq_values = list(sequences.values())

    sequence_batches, index_batches = create_length_batched_sequences(sequences, batch_size)

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

    args = [(batch_seqs, pooling) for batch_seqs in sequence_batches]
    results = worker.embed.starmap(args)

    for embeddings_tensor, index_batch in zip(results, index_batches, strict=True):
        batch_seq_ids = [seq_ids[i] for i in index_batch]
        batch_seq_lengths = [len(seq_values[i]) for i in index_batch]
        yield (embeddings_tensor, batch_seq_ids, batch_seq_lengths)


@app.function(
    gpu=None,
    memory=1024,
    cpu=1.0,
    timeout=24 * HOURS,
)
def embed(
    seqs: dict[str, str],
    model_name: str,
    gpu: str,
    num_gpus: int,
    batch_size: int,
    pooling: str = "mean",
) -> Generator[tuple[torch.Tensor, list[str], list[int]], None, None]:
    """Remote-callable function for protein sequence embedding extraction.

    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)}

    embed_fn = modal.Function.from_name("esm2", "embed")
    for embeddings, seq_ids, seq_lengths in embed_fn.remote_gen(
        seqs=seqs,
        model_name="facebook/esm2_t6_8M_UR50D",
        gpu="T4",
        batch_size=8,
        num_gpus=1,
        pooling="mean",
    ):
        ...
    ```
    """
    yield from calculate_embeddings(seqs, model_name, gpu, num_gpus, batch_size, pooling)
