"""Loading and running the protein masked language model."""
from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache

import spaces  # must be imported before torch (CUDA-related); see app_gradio.py
import pandas as pd
import torch
from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer
from tqdm import tqdm

from src.constants import AMINO_ACIDS

# Sequences per forward call when scoring mutations, not a total batch count:
# total mutations grows with protein length (19 * L), so a *fixed number of
# batches* still lets each batch (and its O(batch * L^2) attention memory)
# grow without bound for long sequences. A fixed batch size keeps peak memory
# roughly constant regardless of protein length.
MUTATION_BATCH_SIZE = 8


@dataclass
class MaskPrediction:
    token: str
    probability: float


@lru_cache(maxsize=None)
def load_model(model_name: str):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer = EsmcTokenizer()
    model = EsmcForMaskedLM.from_pretrained(model_name, device=device)
    model.eval()

    return tokenizer, model


@spaces.GPU
def _sequence_log_probs(sequences: list[str], tokenizer, model) -> torch.Tensor:
    """Unmasked forward pass over full sequences, returned as per-position
    log-probabilities with the leading <cls> token stripped, so that
    `result[i, j]` lines up with `sequences[i][j]`.

    Shape: (len(sequences), len(sequences[0]), vocab_size).
    """
    inputs = tokenizer(sequences, return_tensors="pt", padding=True).to(model.device)
    with torch.no_grad():
        logits = model(**inputs).logits
    log_probs = torch.log_softmax(logits, dim=-1)
    return log_probs[:, 1 : 1 + len(sequences[0]), :]


@spaces.GPU
def _mutation_llrs(
    sequence: str,
    mutations: list[tuple[int, str]],
    wild_type_log_probs: torch.Tensor,
    tokenizer,
    model,
    batch_size: int = MUTATION_BATCH_SIZE,
) -> list[float]:
    """LLR for each `(position, mutant)` pair in `mutations`, processed in
    fixed-size batches of full mutated sequences.

    Only the single log-probability needed per mutant (at its own mutated
    position) is kept — the rest of that sequence's `(L, vocab)` log-softmax
    output is dropped immediately after each batch, instead of being held
    for every mutation at once.
    """
    llrs: list[float] = []
    batch_starts = range(0, len(mutations), batch_size)
    for start in tqdm(batch_starts, desc="Scoring mutations", unit="batch", leave=False):
        chunk = mutations[start : start + batch_size]
        mutated_sequences = [
            sequence[:position] + mutant + sequence[position + 1 :]
            for position, mutant in chunk
        ]
        chunk_log_probs = _sequence_log_probs(mutated_sequences, tokenizer, model)

        for (position, mutant), mutant_log_probs in zip(chunk, chunk_log_probs):
            wild_type_id = tokenizer.convert_tokens_to_ids(sequence[position])
            mutant_id = tokenizer.convert_tokens_to_ids(mutant)
            wild_type_log_prob = wild_type_log_probs[position, wild_type_id]
            mutant_log_prob = mutant_log_probs[position, mutant_id]
            llrs.append(float(mutant_log_prob - wild_type_log_prob))

        del chunk_log_probs
        if model.device.type == "cuda":
            torch.cuda.empty_cache()

    return llrs


def calculate_llr_matrix(sequence: str, model_name: str) -> pd.DataFrame:
    """LLR of every possible single-point mutation at every position.

    Mutant sequences are scored `MUTATION_BATCH_SIZE` at a time (see
    `_mutation_llrs`) rather than all at once, so peak memory stays bounded
    regardless of protein length.

    Returns a `len(AMINO_ACIDS) x len(sequence)` DataFrame where
    `df.loc[aa, col]` is the log-likelihood ratio of mutating the residue at
    `col` to `aa` (mutant-marginal method: the model's own log-probability
    for the mutant residue when reading the full mutated sequence, minus its
    log-probability for the wild-type residue when reading the original
    sequence). The wild-type residue at each position is 0 by definition, so
    it is filled in directly rather than re-scored.
    """
    tokenizer, model = load_model(model_name)
    columns = [f"{residue}{i + 1}" for i, residue in enumerate(sequence)]
    matrix = pd.DataFrame(0.0, index=AMINO_ACIDS, columns=columns)

    wild_type_log_probs = _sequence_log_probs([sequence], tokenizer, model)[0]

    mutations = [
        (position, mutant)
        for position, wild_type in enumerate(sequence)
        for mutant in AMINO_ACIDS
        if mutant != wild_type
    ]
    llrs = _mutation_llrs(sequence, mutations, wild_type_log_probs, tokenizer, model)

    for (position, mutant), llr in zip(mutations, llrs):
        matrix.loc[mutant, matrix.columns[position]] = llr

    return matrix


def predict_masked_position(
    sequence: str,
    position: int,
    model_name: str,
    top_k: int = 5,
) -> list[MaskPrediction]:
    """Mask `sequence[position]` (0-indexed) and return the model's top_k
    amino acid predictions for that position, most probable first."""
    tokenizer, model = load_model(model_name)

    residues = list(sequence)
    residues[position] = tokenizer.mask_token
    masked_sequence = "".join(residues)

    inputs = tokenizer(masked_sequence, return_tensors="pt").to(model.device)
    mask_token_index = torch.where(inputs["input_ids"][0] == tokenizer.mask_token_id)[0]
    if len(mask_token_index) == 0:
        raise ValueError("Mask token not found in tokenized sequence")

    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits[0, mask_token_index[0]], dim=-1)
    top_probs, top_ids = torch.topk(probs, top_k)

    return [
        MaskPrediction(token=tokenizer.decode([tid]).strip(), probability=float(p))
        for tid, p in zip(top_ids, top_probs)
    ]