"""
Loads the SERAPH model once at import time and exposes the prediction
and preset-change handlers used by the Gradio UI.
"""
import spaces
import torch

from config import IDX_TO_LABEL
from dataset import PROTEIN_MAP
from html_builder import build_result_html
from model import load_model

model, tokenizer = load_model()


@spaces.GPU
def predict_structure(sequence_input: str, selected_protein_name: str) -> tuple:
    """Runs ESM-2 + BiLSTM inference and builds the HTML visualization."""
    sequence = sequence_input.upper().strip()
    # Sanitize inputs to amino acid alphabet
    valid_aas = set("ACDEFGHIKLMNPQRSTVWY")
    sequence = "".join([c for c in sequence if c in valid_aas])

    if not sequence:
        empty_html = """
        <div style='padding: 24px; text-align: center; background: rgba(255,255,255,0.02); border: 1px solid rgba(255,255,255,0.08); border-radius: 16px;'>
            <p style='color: #8e8e93; font-family: "JetBrains Mono", monospace;'>Please enter a valid amino acid sequence (e.g., M K W V T F I S L L L L F S S A)...</p>
        </div>
        """
        return empty_html, "", ""

    # Model inference
    tokens = tokenizer(sequence, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        output = model(input_ids=tokens["input_ids"], attention_mask=tokens["attention_mask"])
        preds = output.argmax(dim=-1)[0]

    # Extract predicted structure labels
    pred_labels = [IDX_TO_LABEL[p.item()] for p in preds[1:-1]]
    # Handle truncation alignment safely
    pred_str = "".join(pred_labels[:len(sequence)])
    if len(pred_str) < len(sequence):
        pred_str += "C" * (len(sequence) - len(pred_str))

    # Match with dataset ground truth if selected
    meta = PROTEIN_MAP.get(selected_protein_name, None)
    has_truth = False
    true_ss = ""
    description = ""
    fun_fact = ""

    if meta and meta["sequence"].upper() == sequence:
        has_truth = True
        true_ss = meta["true_ss"]
        description = meta["description"]
        fun_fact = meta["fun_fact"]

    # Calculate Q3 Metric if Ground Truth matches
    q3_score = None
    if has_truth and len(true_ss) == len(sequence):
        matches = sum(1 for i in range(len(sequence)) if pred_str[i] == true_ss[i])
        q3_score = (matches / len(sequence)) * 100.0

    html_output = build_result_html(sequence, pred_str, true_ss if has_truth else None, q3_score, description, fun_fact)

    return html_output, sequence, pred_str


def on_preset_change(selected_name):
    """Populates the sequence textbox when a preset dropdown item is chosen."""
    if selected_name in PROTEIN_MAP:
        item = PROTEIN_MAP[selected_name]
        return item["sequence"]
    return ""