"""Shared ProtT5-XL loading / feature-extraction / h5 I/O helpers.

Used by both app.py (Gradio callbacks, wrapped in @spaces.GPU) and
preprocess/extract_example_feat.py (offline pre-computation of the built-in
example sequences' features into assets/). Keeping this logic in one place
means the cached features in assets/ are guaranteed to match what the app
would compute on the fly.
"""

import re
from pathlib import Path

import h5py
import numpy as np
import torch

T5_REPO_ID = "Rostlab/prot_t5_xl_half_uniref50-enc"

_t5_tokenizer = None
_t5_model = None


def get_device() -> str:
    return "cuda" if torch.cuda.is_available() else "cpu"


def preload_t5_cpu():
    """Deserialize tokenizer + model weights onto CPU, outside any @spaces.GPU
    context.

    This is the slow step (reading ~2.9GB of weights off disk into RAM) and
    it does not need a GPU. Calling it eagerly at app startup means the first
    @spaces.GPU-decorated call only has to do a device transfer (fast), not a
    full from_pretrained() load, so it no longer risks running past the
    ZeroGPU time budget.
    """
    global _t5_tokenizer, _t5_model
    if _t5_model is not None:
        return

    from transformers import AutoTokenizer, T5EncoderModel
    import transformers.utils.import_utils as _hf_utils
    import transformers.modeling_utils as _modeling_utils

    # Rostlab/prot_t5_xl_half_uniref50-enc is only available as .bin (no
    # safetensors). transformers 5.x blocks torch.load on torch < 2.6 due
    # to CVE-2025-32434. We bypass that gate for this specific trusted
    # checkpoint from the official HuggingFace Hub. The check lives in
    # two places — import_utils AND the locally-imported name in
    # modeling_utils — so both must be patched.
    _noop = lambda: None
    _orig_hf = _hf_utils.check_torch_load_is_safe
    _orig_mdl = _modeling_utils.check_torch_load_is_safe

    _hf_utils.check_torch_load_is_safe = _noop
    _modeling_utils.check_torch_load_is_safe = _noop

    try:
        repo_id = T5_REPO_ID
        _t5_tokenizer = AutoTokenizer.from_pretrained(
            repo_id, do_lower_case=False, local_files_only=True
        )
        _t5_model = (
            T5EncoderModel.from_pretrained(
                repo_id, torch_dtype=torch.float32, local_files_only=True
            )
            .eval()
        )
        _t5_model.requires_grad_(False)
    finally:
        _hf_utils.check_torch_load_is_safe = _orig_hf
        _modeling_utils.check_torch_load_is_safe = _orig_mdl


def load_t5():
    """Return (tokenizer, model) with the model on the current device.

    Assumes preload_t5_cpu() has already been called (at app startup). The
    only work left here is moving the already-deserialized model to the GPU
    on the first @spaces.GPU call — a fast operation compared to loading it
    from disk.
    """
    global _t5_model
    preload_t5_cpu()

    dev = get_device()
    if next(_t5_model.parameters()).device.type != dev:
        if dev == "cuda":
            _t5_model = _t5_model.half().to(dev).eval()
        else:
            _t5_model = _t5_model.float().to(dev).eval()

    return _t5_tokenizer, _t5_model


def extract_t5_feature(sequence: str) -> np.ndarray:
    """
    Return mean-pooled ProtT5-XL embedding for a single sequence.
    Output shape: (1024,) — matches the feature dimension LLPSense was trained on.
    """
    tok, mdl = load_t5()
    dev = get_device()

    # Replace ambiguous residues with X (same as original pipeline)
    clean = re.sub(r"[UZOB]", "X", sequence.strip().upper())
    spaced = " ".join(list(clean))

    enc = tok([spaced], add_special_tokens=True, padding="longest", return_tensors="pt")
    input_ids = enc["input_ids"].to(dev)
    attention_mask = enc["attention_mask"].to(dev)

    with torch.no_grad():
        hidden = mdl(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        # Mean-pool over sequence positions (ignoring padding tokens)
        feat = (hidden * attention_mask[..., None]).sum(dim=1) / \
               attention_mask.sum(dim=1, keepdim=True)

    return feat[0].float().cpu().numpy()  # (1024,)


# ── h5 cache I/O ──────────────────────────────────────────────────────────────
FEATURE_TAG = "protein_feat"


def write_feature_h5(filepath, feat: np.ndarray, tag: str = FEATURE_TAG) -> None:
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(filepath, "w") as f:
        f.create_dataset(tag, data=feat, dtype="f", compression="gzip")


def read_feature_h5(filepath, tag: str = FEATURE_TAG) -> np.ndarray:
    with h5py.File(filepath, "r") as f:
        return f[tag][:]