staplebridge/integrations/peptiverse.py

from __future__ import annotations

import math
import os
import sys
import types
from pathlib import Path
from typing import Any, Sequence

import torch


PROPERTY_ALIASES = {
    "hemolysis": "hemolysis",
    "nf": "nf",
    "non_fouling": "nf",
    "non-fouling": "nf",
    "solubility": "solubility",
    "toxicity": "toxicity",
    "permeability": "permeability_penetrance",
    "permeability_penetrance": "permeability_penetrance",
    "pampa": "permeability_pampa",
    "permeability_pampa": "permeability_pampa",
    "caco2": "permeability_caco2",
    "permeability_caco2": "permeability_caco2",
    "half_life": "halflife",
    "halflife": "halflife",
    "binding_affinity": "binding_affinity",
}


def _to_finite_float(x: Any) -> float | None:
    try:
        v = float(x)
    except (TypeError, ValueError):
        return None
    if math.isnan(v) or math.isinf(v):
        return None
    return v


def extract_numeric_score(raw: Any, property_key: str) -> float:
    """Pull the canonical numeric score for ``property_key`` out of ``raw``.

    For ``binding_affinity`` we prefer ``raw["affinity"]`` then ``raw["score"]``.
    For other properties we prefer ``raw["score"]``. Falls back to walking
    other plausible keys if a top-level dict is given. Raises ``ValueError``
    if no finite float can be produced. When the priority keys do not yield a
    value but the fallback dict-walk does, a ``warnings.warn`` is emitted so
    silent mis-extraction is surfaced.
    """
    import warnings as _warnings

    pkey = (property_key or "").lower()
    candidates: list[Any] = []
    if isinstance(raw, dict):
        if pkey == "binding_affinity":
            candidates += [raw.get("affinity"), raw.get("score")]
        else:
            candidates += [raw.get("score"), raw.get(pkey)]
        candidates += [raw.get("value"), raw.get("prediction"), raw.get("mean")]
    elif isinstance(raw, (list, tuple)) and raw:
        candidates.append(raw[0])
    else:
        candidates.append(raw)

    for c in candidates:
        v = _to_finite_float(c)
        if v is not None:
            return v
    if isinstance(raw, dict):
        finite_pairs = [(k, _to_finite_float(v)) for k, v in raw.items()]
        finite_pairs = [(k, v) for k, v in finite_pairs if v is not None]
        if finite_pairs:
            if len(finite_pairs) > 1:
                _warnings.warn(
                    f"PeptiVerse score extraction for property={property_key!r}: "
                    f"no priority key matched; falling back to dict walk over "
                    f"{[k for k, _ in finite_pairs]} -> picked first finite value "
                    f"{finite_pairs[0][1]} from key {finite_pairs[0][0]!r}. "
                    "This may be wrong — verify the predictor output schema.",
                    stacklevel=2,
                )
            return finite_pairs[0][1]
    raise ValueError(
        f"Cannot extract finite numeric score for property={property_key!r} from raw={raw!r}"
    )


class PeptiVerseUnavailable(RuntimeError):
    pass


class PeptiVerseWrapper:
    """Local-import wrapper around ChatterjeeLab/PeptiVerse `inference.py`.

    The HF repo is expected to be cloned to disk at ``peptiverse_root``; we
    insert that path into ``sys.path`` and import ``inference.PeptiVersePredictor``
    directly. No remote HTTP API is used.

    Parameters
    ----------
    peptiverse_root : path to the local PeptiVerse repo clone.
    manifest_path   : best_models.txt path (defaults to ``peptiverse_root/best_models.txt``).
    classifier_weight_root : weight root (defaults to ``peptiverse_root``).
    device          : ``"cpu"`` or ``"cuda"``.
    strict          : if True, raise on any init error.
    uncertainty     : default ``uncertainty`` flag passed to the underlying
                      predictor unless overridden per-call.
    cache_enabled   : memoize identical (property, mode, input, target, uncertainty)
                      calls to avoid re-running the underlying torch models.
    """

    def __init__(
        self,
        peptiverse_root: str | Path,
        manifest_path: str | Path | None = None,
        classifier_weight_root: str | Path | None = None,
        device: str | None = None,
        strict: bool = False,
        uncertainty: bool = False,
        cache_enabled: bool = True,
        hf_cache_dir: str | Path | None = None,
        esm_model_name_or_path: str | Path | None = None,
        peptideclm_model_name_or_path: str | Path | None = None,
        chemberta_model_name_or_path: str | Path | None = None,
        offline: bool = False,
        batch_size: int = 32,
    ) -> None:
        self.root = Path(peptiverse_root) if peptiverse_root else None
        self.manifest_path = Path(manifest_path) if manifest_path else None
        self.classifier_weight_root = (
            Path(classifier_weight_root) if classifier_weight_root else None
        )
        self.device = device
        self.strict = strict
        self.uncertainty = bool(uncertainty)
        self.cache_enabled = bool(cache_enabled)
        self.hf_cache_dir = Path(hf_cache_dir) if hf_cache_dir else None
        self.esm_model_name_or_path = (
            Path(esm_model_name_or_path) if esm_model_name_or_path else None
        )
        self.peptideclm_model_name_or_path = (
            Path(peptideclm_model_name_or_path)
            if peptideclm_model_name_or_path
            else None
        )
        self.chemberta_model_name_or_path = (
            Path(chemberta_model_name_or_path)
            if chemberta_model_name_or_path
            else None
        )
        self.offline = bool(offline)
        self.predictor = None
        self.available = False
        self.init_error: str | None = None
        self._cache: dict[tuple, dict[str, Any]] = {}
        self.cache_hits = 0
        self.cache_misses = 0
        self.batch_size = max(1, int(batch_size))
        self.batch_chunks = 0
        self.batch_oom_retries = 0
        self.batch_oom_giveups = 0
        self._init_predictor()

    @staticmethod
    def _snapshot_from_cache(cache_root: Path, model_id: str) -> Path | None:
        """Resolve a cached Hugging Face repo to a concrete local snapshot."""
        repo = cache_root / "hub" / f"models--{model_id.replace('/', '--')}"
        snapshots = repo / "snapshots"
        ref = repo / "refs" / "main"
        if ref.is_file():
            revision = ref.read_text(encoding="utf-8").strip()
            candidate = snapshots / revision
            if candidate.is_dir():
                return candidate.resolve()
        candidates = sorted(path for path in snapshots.glob("*") if path.is_dir())
        if len(candidates) == 1:
            return candidates[0].resolve()
        return None

    def _local_hf_models(self) -> dict[str, Path]:
        configured = {
            "esm_name": self.esm_model_name_or_path,
            "clm_name": self.peptideclm_model_name_or_path,
            "chemberta_name": self.chemberta_model_name_or_path,
        }
        model_ids = {
            "esm_name": "facebook/esm2_t33_650M_UR50D",
            "clm_name": "aaronfeller/PeptideCLM-23M-all",
            "chemberta_name": "DeepChem/ChemBERTa-77M-MLM",
        }
        resolved: dict[str, Path] = {}
        for argument, configured_path in configured.items():
            path = configured_path
            if path is None and self.hf_cache_dir is not None:
                path = self._snapshot_from_cache(
                    self.hf_cache_dir, model_ids[argument]
                )
            if path is not None:
                path = path.expanduser().resolve()
                if not path.is_dir():
                    raise FileNotFoundError(
                        f"local PeptiVerse HF model directory missing: {path}"
                    )
                resolved[argument] = path
        return resolved

    def _init_predictor(self) -> None:
        if self.hf_cache_dir is not None:
            cache_root = self.hf_cache_dir.expanduser().resolve()
            # Deliberately override inherited values.  A stale user-level HF
            # cache was the source of tokenizer lookup failures in offline
            # runs; this wrapper's configured local model root is authoritative.
            os.environ["HF_HOME"] = str(cache_root)
            os.environ["HUGGINGFACE_HUB_CACHE"] = str(cache_root / "hub")
            os.environ["TRANSFORMERS_CACHE"] = str(cache_root / "hub")
        if self.offline:
            os.environ["HF_HUB_OFFLINE"] = "1"
            os.environ["TRANSFORMERS_OFFLINE"] = "1"

        if self.root is None or not self.root.exists():
            self.init_error = f"peptiverse_root does not exist: {self.root}"
            if self.strict:
                raise PeptiVerseUnavailable(self.init_error)
            return
        if not (self.root / "inference.py").exists():
            self.init_error = f"inference.py not found under {self.root}"
            if self.strict:
                raise PeptiVerseUnavailable(self.init_error)
            return

        sys.path.insert(0, str(self.root))
        try:
            # PeptiVerse imports ``tokenizer.my_tokenizers`` from its own
            # source tree. Some environments also install an unrelated
            # top-level ``tokenizer`` package; once that regular package is in
            # sys.modules it shadows PeptiVerse's namespace directory even
            # though the PeptiVerse root is first on sys.path. Bind the local
            # namespace explicitly before importing inference.
            local_tokenizer = self.root / "tokenizer"
            if (local_tokenizer / "my_tokenizers.py").is_file():
                loaded = sys.modules.get("tokenizer")
                loaded_paths = [
                    str(Path(value).resolve())
                    for value in (getattr(loaded, "__path__", None) or [])
                ]
                if str(local_tokenizer.resolve()) not in loaded_paths:
                    for name in list(sys.modules):
                        if name == "tokenizer" or name.startswith("tokenizer."):
                            del sys.modules[name]
                    namespace = types.ModuleType("tokenizer")
                    namespace.__path__ = [str(local_tokenizer)]  # type: ignore[attr-defined]
                    namespace.__package__ = "tokenizer"
                    sys.modules["tokenizer"] = namespace
            from inference import PeptiVersePredictor  # type: ignore

            # Several PeptiVerse MAPIE calibration artifacts were serialized
            # by a training script where this estimator lived in __main__.
            # Re-export the identical local inference.py class there before
            # joblib loads those artifacts. Point predictions do not use MAPIE
            # when uncertainty=False, but resolving the class avoids noisy
            # load failures and preserves optional uncertainty availability.
            try:
                from inference import PassthroughRegressor  # type: ignore

                main_module = sys.modules.get("__main__")
                if main_module is not None and not hasattr(
                    main_module, "PassthroughRegressor"
                ):
                    setattr(
                        main_module,
                        "PassthroughRegressor",
                        PassthroughRegressor,
                    )
            except ImportError:
                pass

            manifest = self.manifest_path or (self.root / "best_models.txt")
            weight_root = self.classifier_weight_root or self.root
            kwargs: dict[str, Any] = {
                "manifest_path": str(manifest),
                "classifier_weight_root": str(weight_root),
            }
            if self.device is not None:
                kwargs["device"] = self.device
            for argument, model_path in self._local_hf_models().items():
                kwargs[argument] = str(model_path)
            self.predictor = PeptiVersePredictor(**kwargs)
            self.available = True
        except Exception as exc:  # noqa: BLE001 - surface any import/init failure
            self.init_error = f"{type(exc).__name__}: {exc}"
            self.available = False
            if self.strict:
                raise PeptiVerseUnavailable(self.init_error) from exc

    @staticmethod
    def normalize_property_name(key: str) -> str:
        return PROPERTY_ALIASES.get(key.lower(), key)

    @staticmethod
    def extract_numeric_score(raw: Any, property_key: str) -> float:
        return extract_numeric_score(raw, property_key)

    def _ensure_available(self) -> None:
        if not self.available or self.predictor is None:
            raise PeptiVerseUnavailable(
                self.init_error or "PeptiVerse predictor is not available"
            )

    def clear_cache(self) -> None:
        self._cache.clear()
        self.cache_hits = 0
        self.cache_misses = 0

    def _cache_key(
        self, property_key: str, mode: str, input_str: str, target_seq: str, uncertainty: bool
    ) -> tuple:
        return (property_key, mode, input_str, target_seq, bool(uncertainty))

    def predict_property(
        self,
        property_key: str,
        input_str: str,
        mode: str = "wt",
        uncertainty: bool | None = None,
    ) -> dict[str, Any]:
        """Run a single non-binding property head.

        Returns ``{"backend_used": "peptiverse", "raw": <dict>, "score": float,
        "property": <std_key>, "mode": <mode>, "cached": bool}``.
        """
        self._ensure_available()
        std_key = self.normalize_property_name(property_key)
        unc = self.uncertainty if uncertainty is None else bool(uncertainty)
        key = self._cache_key(std_key, mode, input_str, "", unc)
        if self.cache_enabled and key in self._cache:
            self.cache_hits += 1
            cached = dict(self._cache[key])
            cached["cached"] = True
            return cached

        raw = self.predictor.predict_property(  # type: ignore[union-attr]
            std_key, mode, input_str, unc
        )
        score = extract_numeric_score(raw, std_key)
        out = {
            "backend_used": "peptiverse",
            "property": std_key,
            "mode": mode,
            "raw": raw,
            "score": float(score),
            "cached": False,
        }
        if self.cache_enabled:
            self._cache[key] = {k: v for k, v in out.items() if k != "cached"}
        self.cache_misses += 1
        return out

    def predict_property_batch(
        self,
        property_key: str,
        input_strs: Sequence[str],
        mode: str = "wt",
        uncertainty: bool | None = None,
        batch_size: int | None = None,
    ) -> list[dict[str, Any]]:
        """Run one non-binding property head over many inputs.

        Same per-item contract as :meth:`predict_property`, returned in input
        order. Duplicates and cache hits are computed once. The speedup comes
        from batching the *embedding* forward pass; the heads still see one
        input at a time, so each score is produced by exactly the code path
        :meth:`predict_property` would have used.

        Batching an embedder changes the padded sequence length, which perturbs
        the pooled embedding at ~1e-7. That is why only the embeddings are
        batched here and the results are cached under the same keys as the
        scalar path: a later scalar call returns the batch-computed value, so a
        run cannot mix the two conventions for the same input. Verify numeric
        agreement for any new property head before relying on it.

        On CUDA OOM the offending chunk is retried at half the batch size, down
        to scalar, so a long tail of large molecules degrades rather than fails.
        """
        self._ensure_available()
        std_key = self.normalize_property_name(property_key)
        if std_key == "binding_affinity":
            raise ValueError("use predict_binding_affinity for binding affinity")
        unc = self.uncertainty if uncertainty is None else bool(uncertainty)
        inputs = list(input_strs)
        if not inputs:
            return []

        # Deduplicate while preserving first-seen order; only uncached, unique
        # inputs reach the backend.
        pending: list[str] = []
        seen: set[str] = set()
        for value in inputs:
            key = self._cache_key(std_key, mode, value, "", unc)
            if self.cache_enabled and key in self._cache:
                continue
            if value in seen:
                continue
            seen.add(value)
            pending.append(value)

        if pending:
            self._embed_batch(std_key, mode, pending, unc, batch_size)
        # Every input now either was cached or has a warm embedding, so this
        # loop is the unmodified scalar path and defines the returned values.
        return [
            self.predict_property(std_key, value, mode=mode, uncertainty=unc)
            for value in inputs
        ]

    def _embed_batch(
        self,
        std_key: str,
        mode: str,
        pending: Sequence[str],
        uncertainty: bool,
        batch_size: int | None,
    ) -> None:
        """Warm the underlying embedder's cache for ``pending`` in batches."""
        embedder, pooled_kind = self._embedder_for(std_key, mode)
        if embedder is None:
            return
        size = int(batch_size or self.batch_size)
        index = 0
        items = list(pending)
        while index < len(items):
            chunk = items[index : index + size]
            try:
                self._embed_chunk(embedder, pooled_kind, chunk)
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                if len(chunk) == 1:
                    # Nothing left to split; let the scalar path surface it.
                    self.batch_oom_giveups += 1
                    index += 1
                    continue
                size = max(1, len(chunk) // 2)
                self.batch_oom_retries += 1
                continue
            index += len(chunk)
            self.batch_chunks += 1

    def _embedder_for(self, std_key: str, mode: str) -> tuple[Any, str]:
        """The embedder and pooling kind the head for ``std_key`` will request."""
        predictor = self.predictor
        meta = getattr(predictor, "meta", {}).get((std_key, mode))
        if meta is None:
            return None, ""
        try:
            embedder = predictor._get_embedder(meta["emb_tag"])  # type: ignore[union-attr]
        except (KeyError, ValueError):
            return None, ""
        # torch_ckpt heads consume unpooled token embeddings; everything else
        # consumes the masked mean pool. Mirrors PeptiVersePredictor._get_features.
        return embedder, ("unpooled" if meta.get("kind") == "torch_ckpt" else "pooled")

    @staticmethod
    def _embed_chunk(embedder: Any, pooled_kind: str, chunk: Sequence[str]) -> None:
        """Forward ``chunk`` once and store per-item results in the embedder cache.

        Writes into the embedder's own ``_cache_pooled``/``_cache_unpooled`` so
        the scalar ``pooled``/``unpooled`` calls that follow find them. Items
        already cached are skipped.
        """
        cache_attr = "_cache_pooled" if pooled_kind == "pooled" else "_cache_unpooled"
        cache = getattr(embedder, cache_attr, None)
        if cache is None:
            return
        stripped = [str(value).strip() for value in chunk]
        fresh = [value for value in dict.fromkeys(stripped) if value not in cache]
        if not fresh:
            return

        tokenized = embedder._tokenize(fresh)
        with torch.no_grad():
            hidden = embedder.model(
                input_ids=tokenized["input_ids"],
                attention_mask=tokenized["attention_mask"],
            ).last_hidden_state
        valid = embedder._valid_mask(tokenized["input_ids"], tokenized["attention_mask"])

        if pooled_kind == "pooled":
            weights = valid.unsqueeze(-1).float()
            pooled = (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-9)
            for position, value in enumerate(fresh):
                cache[value] = pooled[position : position + 1]
            return

        for position, value in enumerate(fresh):
            row = valid[position]
            features = hidden[position : position + 1, row, :]
            mask = torch.ones(
                (1, features.shape[1]), dtype=torch.bool, device=features.device
            )
            cache[value] = (features, mask)

    def predict_binding_affinity(
        self,
        binder_str: str,
        target_seq: str,
        mode: str = "wt",
        uncertainty: bool | None = None,
    ) -> dict[str, Any]:
        """Run the binding-affinity head against a specific protein target."""
        self._ensure_available()
        if not target_seq:
            raise ValueError("predict_binding_affinity requires target_seq")
        unc = self.uncertainty if uncertainty is None else bool(uncertainty)
        key = self._cache_key("binding_affinity", mode, binder_str, target_seq, unc)
        if self.cache_enabled and key in self._cache:
            self.cache_hits += 1
            cached = dict(self._cache[key])
            cached["cached"] = True
            return cached

        raw = self.predictor.predict_binding_affinity(  # type: ignore[union-attr]
            mode, target_seq, binder_str, unc
        )
        score = extract_numeric_score(raw, "binding_affinity")
        out = {
            "backend_used": "peptiverse",
            "property": "binding_affinity",
            "mode": mode,
            "raw": raw,
            "score": float(score),
            "cached": False,
        }
        if self.cache_enabled:
            self._cache[key] = {k: v for k, v in out.items() if k != "cached"}
        self.cache_misses += 1
        return out