from __future__ import annotations

import modal
import torch

from apb.biomodal.esm2.cache_model import app as cache_app
from apb.biomodal.esm2.constants import (
    APP_NAME,
    HOURS,
    MAX_SEQ_LENGTH,
    MODEL_STORAGE_PATH,
    console,
    image,
    volume,
)
from apb.biomodal.esm2.ops import calculate_pseudo_log_likelihood

with image.imports():
    from transformers import AutoTokenizer, EsmForMaskedLM

app = modal.App(name=APP_NAME, image=image)
app.include(cache_app)


@app.cls(
    volumes={MODEL_STORAGE_PATH: volume},
    gpu="any",
    max_containers=1,
    timeout=24 * HOURS,
    scaledown_window=60,
)
class EsmModel:
    """A Modal class for running ESM models on GPU.

    Supports pseudo-log-likelihood calculation and embedding extraction.
    """

    def __init__(self, model_name: str):
        self.model_name = model_name

    @modal.enter()
    def load_model(self):
        assert torch.cuda.is_available()

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            cache_dir=MODEL_STORAGE_PATH,
            local_files_only=True,
        )
        self.model = EsmForMaskedLM.from_pretrained(
            self.model_name,
            cache_dir=MODEL_STORAGE_PATH,
            local_files_only=True,
        )
        self.model.to("cuda")  # type: ignore
        self.model.eval()

    @modal.method()
    def pseudo_log_likelihood(self, sequences: list[str]) -> torch.Tensor:
        tokenizer_output = self.tokenizer(
            sequences,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=MAX_SEQ_LENGTH + 2,
        )

        token_ids = tokenizer_output["input_ids"]
        attention_mask = tokenizer_output["attention_mask"]

        CLS = self.tokenizer.cls_token_id
        EOS = self.tokenizer.eos_token_id
        PAD = self.tokenizer.pad_token_id
        token_mask = (token_ids != CLS) & (token_ids != EOS) & (token_ids != PAD)

        plls = calculate_pseudo_log_likelihood(
            token_ids=token_ids,
            attention_mask=attention_mask,
            token_mask=token_mask,
            model=self.model,
            mask_token_id=self.tokenizer.mask_token_id,
            quiet=False,
        ).cpu()

        console.print("Finished batch.")

        return plls

    @modal.method()
    def embed(self, sequences: list[str], pooling: str = "mean") -> torch.Tensor:
        tokenizer_output = self.tokenizer(
            sequences,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=MAX_SEQ_LENGTH + 2,
        )

        input_ids = tokenizer_output["input_ids"].to("cuda")
        attention_mask = tokenizer_output["attention_mask"].to("cuda")

        with torch.no_grad():
            outputs = self.model.esm(input_ids, attention_mask=attention_mask)
            last_hidden_state = outputs.last_hidden_state.cpu()

        if pooling == "mean":
            embeddings = torch.stack(
                [
                    last_hidden_state[i, 1 : min(MAX_SEQ_LENGTH, len(seq)) + 1, :].mean(0)
                    for i, seq in enumerate(sequences)
                ]
            )
            console.print("Finished embedding batch (mean pooling).")
            return embeddings

        console.print("Finished embedding batch (no pooling).")
        return last_hidden_state
