from pathlib import Path
from typing import Annotated

import modal
import torch
import typer
from biotite.sequence.io.fasta import FastaFile

from apb.biomodal.esm2.constants import MAX_SEQ_LENGTH
from apb.biomodal.utils import validate_modal

cli = typer.Typer(pretty_exceptions_enable=False, add_completion=False)


@cli.callback()
def _callback() -> None:
    validate_modal()


def _load_sequences(fasta_file: Path) -> dict[str, str]:
    return {defline: seq for defline, seq in FastaFile.read_iter(fasta_file)}


MODEL_HELP = (
    "ESM model name. Options: facebook/esm2_t6_8M_UR50D (8M), "
    "facebook/esm2_t12_35M_UR50D (35M), facebook/esm2_t30_150M_UR50D (150M), "
    "facebook/esm2_t33_650M_UR50D (650M), facebook/esm2_t36_3B_UR50D (3B), "
    "facebook/esm2_t48_15B_UR50D (15B)"
)

GPU_HELP = (
    "GPU type. Options: any (let Modal choose), t4 (16GB), a10g (24GB), a100 (40GB), h100 (80GB)"
)

BATCH_SIZE_HELP = (
    "Number of sequences to process per batch. Lower if you're getting out of memory errors."
)


@cli.command()
def perplexity(
    fasta_file: Annotated[Path, typer.Option(help="Path to input FASTA file containing sequences")],
    output: Annotated[Path, typer.Option(help="Path to output CSV file for perplexity results")],
    model_name: Annotated[str, typer.Option(help=MODEL_HELP)],
    gpu: Annotated[str, typer.Option(help=GPU_HELP)],
    num_gpus: Annotated[
        int, typer.Option(help="Number of GPU workers to spawn for parallel processing.")
    ],
    batch_size: Annotated[int, typer.Option(help=BATCH_SIZE_HELP)],
) -> None:
    from apb.biomodal.esm2.model import app as modal_app
    from apb.biomodal.esm2.perplexity import pseudo_perplexity

    seqs = _load_sequences(fasta_file)

    with modal.enable_output(), modal_app.run():
        pppl = pseudo_perplexity.local(
            seqs,
            model_name=model_name,
            gpu=gpu,
            num_gpus=num_gpus,
            batch_size=batch_size,
        )

    pppl.to_csv(output, index=False)


@cli.command()
def embed(
    fasta_file: Annotated[Path, typer.Option(help="Path to input FASTA file containing sequences")],
    output_dir: Annotated[Path, typer.Option(help="Directory to save embedding .pt files")],
    model_name: Annotated[str, typer.Option(help=MODEL_HELP)],
    gpu: Annotated[str, typer.Option(help=GPU_HELP)],
    num_gpus: Annotated[
        int, typer.Option(help="Number of GPU workers to spawn for parallel processing.")
    ],
    batch_size: Annotated[int, typer.Option(help=BATCH_SIZE_HELP)],
    pooling: Annotated[
        str, typer.Option(help="Pooling strategy. 'mean' for mean-pooled, 'none' for per-residue.")
    ] = "mean",
) -> None:
    from apb.biomodal.esm2.embedding import embed as embed_deployed
    from apb.biomodal.esm2.model import app as modal_app

    seqs = _load_sequences(fasta_file)
    output_dir.mkdir(parents=True, exist_ok=True)

    with modal.enable_output(), modal_app.run():
        for embeddings, seq_ids, seq_lengths in embed_deployed.local(
            seqs,
            model_name=model_name,
            gpu=gpu,
            num_gpus=num_gpus,
            batch_size=batch_size,
            pooling=pooling,
        ):
            for i, seq_id in enumerate(seq_ids):
                if pooling == "mean":
                    tensor = embeddings[i]
                else:
                    end = min(MAX_SEQ_LENGTH, seq_lengths[i]) + 1
                    tensor = embeddings[i, 1:end]
                torch.save(tensor.to(torch.float16), output_dir / f"{seq_id}.pt")


if __name__ == "__main__":
    cli()
