models.py

"""
HuggingFace model integration.
Uses ESM-2 (facebook/esm2_t6_8M_UR50D) for protein sequence embeddings
and structural property prediction. Lazy-loaded on first use.
"""
from __future__ import annotations
import os, time
import numpy as np

_esm_model = None
_esm_tokenizer = None

MODEL_ID = os.getenv("ESM_MODEL", "facebook/esm2_t6_8M_UR50D")


def _load_esm():
    global _esm_model, _esm_tokenizer
    if _esm_model is not None:
        return
    try:
        from transformers import AutoTokenizer, AutoModel
        import torch
        _esm_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        _esm_model = AutoModel.from_pretrained(MODEL_ID)
        _esm_model.eval()
    except Exception as e:
        print(f"[models] ESM-2 load failed: {e}. Falling back to mock embeddings.")


def embed_sequence(sequence: str) -> list[float]:
    """
    Return a 320-dim embedding for the protein sequence.
    Uses ESM-2 if available, otherwise a deterministic hash-based mock.
    """
    if not sequence:
        return [0.0] * 320

    _load_esm()

    if _esm_model is not None:
        try:
            import torch
            seq = sequence[:1022]  # ESM-2 max length
            inputs = _esm_tokenizer(seq, return_tensors="pt", truncation=True)
            with torch.no_grad():
                out = _esm_model(**inputs)
            # Mean pool over residue dimension
            emb = out.last_hidden_state[0, 1:-1].mean(dim=0).numpy().tolist()
            return emb[:320] if len(emb) >= 320 else emb + [0.0] * (320 - len(emb))
        except Exception as e:
            print(f"[models] ESM-2 inference error: {e}")

    # Deterministic fallback: hash-based mock embedding
    import hashlib
    h = hashlib.sha256(sequence.encode()).digest()
    rng = np.random.default_rng(int.from_bytes(h[:8], "big"))
    return rng.uniform(-1, 1, 320).tolist()


def predict_disorder(sequence: str) -> list[float]:
    """
    Predict per-residue disorder probability (0=ordered, 1=disordered).
    Simple heuristic using amino acid physicochemical properties.
    """
    DISORDER_PRONE = set("GSEQKAPRT")
    ORDER_PRONE = set("CFYWILVMH")
    scores = []
    for aa in sequence.upper():
        if aa in DISORDER_PRONE:
            scores.append(0.7)
        elif aa in ORDER_PRONE:
            scores.append(0.2)
        else:
            scores.append(0.45)
    # Smooth with a sliding window
    w = 9
    padded = [scores[0]] * (w // 2) + scores + [scores[-1]] * (w // 2)
    smoothed = [
        sum(padded[i:i + w]) / w for i in range(len(scores))
    ]
    return smoothed


def classify_secondary_structure_heuristic(sequence: str) -> dict:
    """
    Fast heuristic secondary structure estimate.
    Returns fraction of helix, sheet, coil residues.
    """
    HELIX = set("AELMQKR")
    SHEET = set("VIFYW")
    total = len(sequence) or 1
    h = sum(1 for aa in sequence.upper() if aa in HELIX) / total
    s = sum(1 for aa in sequence.upper() if aa in SHEET) / total
    return {
        "helix_fraction": round(h, 3),
        "sheet_fraction": round(s, 3),
        "coil_fraction": round(max(0, 1 - h - s), 3),
    }


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine similarity between two embedding vectors."""
    a, b = np.array(a), np.array(b)
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / norm) if norm > 0 else 0.0