src/pino/embeddings.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
from rdkit import Chem, RDLogger
from rdkit.Chem import AllChem, Descriptors

from .registry import AromaRegistry
from .thermo.geometry import extract_3d_shape_features
from .thermo.naturals import OPENPOM_DIM, SHAPE_DIM, resolve_natural_oil_vectors

logger = logging.getLogger("pino.embeddings")


class RealOpenPOMUnavailableError(RuntimeError):
    """Raised when a caller requires real POM features but no model is loaded."""

STRUCTURAL_DIM = OPENPOM_DIM  # 138 (Morgan fallback structural block)
REAL_POM_DIM = 256  # genuine OpenPOM v1.0.0 penultimate embedding (ablation-proven)
SHAPE_DIM = SHAPE_DIM  # 7
PHYSICAL_DIM = 2  # molecular_weight, logp
# --- morgan_2048_rp arm (round-7 panel response): 2048-bit radius-2 Morgan +
# 10 physicochemical descriptors = 2058-D structural block; full token 2071-D;
# fixed seeded random projection down to the OpenPOM arm's combined width so
# both trained arms carry an IDENTICAL input-projection parameter count
# (parameter-parity control requested by the evaluation reviewer).
MORGAN2048_STRUCTURAL_DIM = 2058
MORGAN2048_TOKEN_DIM = MORGAN2048_STRUCTURAL_DIM + SHAPE_DIM + PHYSICAL_DIM + 4  # 2071
MORGAN2048_RP_SEED = 20260725
FUNCTIONAL_DIM = 4  # ester, alcohol, ketone, aldehyde
COMBINED_DIM = STRUCTURAL_DIM + SHAPE_DIM + PHYSICAL_DIM + FUNCTIONAL_DIM  # 151
COMBINED_DIM_REAL_POM = REAL_POM_DIM + SHAPE_DIM + PHYSICAL_DIM + FUNCTIONAL_DIM  # 269

_REPO_ROOT = Path(__file__).resolve().parents[2]
_POM_DIR = _REPO_ROOT / "artifacts" / "pom"
_DISJOINT_DIR = _REPO_ROOT / "artifacts" / "pom_disjoint"

# Mute RDKit's internal C++ warning streams globally.
rd_logger = RDLogger.logger()
rd_logger.setLevel(RDLogger.CRITICAL)


# UNIFAC subgroup ID mapping to functional families (Dortmund-UNIFAC convention).
# These IDs are stable in the thermo.unifac DOUFSG parameter set.
FUNCTIONAL_GROUP_IDS = {
    "ester": {21, 22, 23},
    "alcohol": {14, 15, 16},
    "ketone": {18, 19},
    "aldehyde": {20},
}


def _functional_group_flags(unifac_groups: dict[str, Any] | str | None) -> np.ndarray:
    """Return a 4-D binary vector for ester, alcohol, ketone, aldehyde."""
    flags = np.zeros(FUNCTIONAL_DIM, dtype=np.float32)
    if not unifac_groups:
        return flags

    if isinstance(unifac_groups, str):
        try:
            unifac_groups = json.loads(unifac_groups)
        except Exception:
            return flags

    group_ids = set()
    for key in unifac_groups:
        try:
            group_ids.add(int(key))
        except ValueError:
            # If keys are textual names (legacy), look for matching subgroup names.
            k = str(key).lower()
            if "coo" in k or "ester" in k:
                flags[0] = 1.0
            if "oh" in k:
                flags[1] = 1.0
            if "ketone" in k or "co" == k:
                flags[2] = 1.0
            if "cho" in k or "aldehyde" in k:
                flags[3] = 1.0
            return flags

    if group_ids & FUNCTIONAL_GROUP_IDS["ester"]:
        flags[0] = 1.0
    if group_ids & FUNCTIONAL_GROUP_IDS["alcohol"]:
        flags[1] = 1.0
    if group_ids & FUNCTIONAL_GROUP_IDS["ketone"]:
        flags[2] = 1.0
    if group_ids & FUNCTIONAL_GROUP_IDS["aldehyde"]:
        flags[3] = 1.0
    return flags


class OlfactoryEmbeddingEngine:
    """
    Write-through cached embedding engine for 151-D olfactory molecule tokens.

    Combines a 138-D OpenPOM-style structural vector, a 7-D 3D shape/chirality
    descriptor, and a 6-D physical-functional payload (molecular weight, LogP,
    ester/alcohol/ketone/aldehyde flags) from the SQLite AromaRegistry.

    Missing 3D features are computed once and written back to the registry.
    Empty or invalid inputs short-circuit to a zero fallback without invoking RDKit.
    """

    def __init__(
        self,
        use_fallback: bool = True,
        global_median_logp: float = 2.15,
        structural_source: str = "morgan",
    ) -> None:
        self.use_fallback = use_fallback
        self.structural_source = structural_source
        self.use_real_pom = structural_source == "openpom_256"
        self.use_morgan_2048_rp = structural_source == "morgan_2048_rp"
        self.use_morgan_2048_direct = structural_source == "morgan_2048_direct"
        self.use_disjoint_pom = structural_source == "disjoint_256"
        self.structural_dim = (
            REAL_POM_DIM if (self.use_real_pom or self.use_disjoint_pom)
            else MORGAN2048_STRUCTURAL_DIM if (self.use_morgan_2048_rp or self.use_morgan_2048_direct)
            else STRUCTURAL_DIM
        )
        self.shape_dim = SHAPE_DIM
        self.physical_dim = PHYSICAL_DIM
        self.functional_dim = FUNCTIONAL_DIM
        self.embedding_dim = (
            COMBINED_DIM_REAL_POM if (self.use_real_pom or self.use_morgan_2048_rp or self.use_disjoint_pom)
            else MORGAN2048_TOKEN_DIM if self.use_morgan_2048_direct
            else COMBINED_DIM
        )
        self._rp_matrix: np.ndarray | None = None
        self.registry = AromaRegistry()
        self._pom_index: dict[str, int] = {}
        self._pom_matrix: np.ndarray | None = None
        if self.use_real_pom:
            self._load_real_pom()
        if self.use_disjoint_pom:
            self._load_disjoint_pom()
        self.global_median_logp = self._compute_median_logp() if global_median_logp is None else global_median_logp
        if not self.use_fallback:
            self._initialize_real_openpom()

    def _load_real_pom(self) -> None:
        """Load the genuine OpenPOM v1.0.0 ensemble (256-dim, 10-model mean).

        Indexed by CAS (with some SMILES:-prefixed keys). Fail closed if the
        provenance-manifest asset is absent so a Morgan feature is never
        silently substituted for POM."""
        emb_path = _POM_DIR / "pom_embeddings.npy"
        idx_path = _POM_DIR / "pom_index.json"
        if not (emb_path.exists() and idx_path.exists()):
            raise RealOpenPOMUnavailableError(
                f"Genuine POM asset missing under {_POM_DIR}; cannot honour structural_source='openpom_256'"
            )
        self._pom_matrix = np.load(emb_path).astype(np.float32)
        idx = json.loads(idx_path.read_text())
        ids = idx.get("ids", [])
        self._pom_index = {self._normalized_cas(k): i for i, k in enumerate(ids)}

    def _load_disjoint_pom(self) -> None:
        """Load the benchmark-disjoint MPNN encoder embeddings (256-dim).

        Same CAS-indexed layout as the genuine-POM asset, produced by
        scripts/build_disjoint_asset_and_gate.py from the round-7 disjoint
        encoder (mattbitzesty/pino-disjoint-encoder). Fail closed if absent:
        a Morgan feature must never be silently substituted for the disjoint
        arm's representation. Reuses _pom_index/_pom_matrix slots (the two
        arms are never active simultaneously)."""
        emb_path = _DISJOINT_DIR / "disjoint_embeddings.npy"
        idx_path = _DISJOINT_DIR / "disjoint_index.json"
        if not (emb_path.exists() and idx_path.exists()):
            raise RealOpenPOMUnavailableError(
                f"Disjoint-encoder asset missing under {_DISJOINT_DIR}; cannot honour structural_source='disjoint_256'"
            )
        self._pom_matrix = np.load(emb_path).astype(np.float32)
        idx = json.loads(idx_path.read_text())
        ids = idx.get("ids", [])
        self._pom_index = {self._normalized_cas(k): i for i, k in enumerate(ids)}

    def _compute_median_logp(self) -> float:
        rows = self.registry._conn.execute(
            "SELECT logp FROM aroma_chemicals WHERE logp IS NOT NULL AND logp != 0"
        ).fetchall()
        if not rows:
            return 2.15
        return float(np.median([r[0] for r in rows]))

    def _initialize_real_openpom(self) -> None:
        # Fail closed.  Silently substituting Morgan features here makes a
        # Morgan-vs-POM ablation compare the same representation under two names.
        raise RealOpenPOMUnavailableError(
            "Real OpenPOM model is not configured; use_fallback=False cannot be satisfied"
        )

    def _normalized_cas(self, cas: str | None) -> str:
        return str(cas).replace("NATURAL:", "").strip() if cas else ""

    def _is_valid_smiles(self, smiles: str) -> bool:
        if not smiles or not isinstance(smiles, str):
            return False
        s = smiles.strip()
        return bool(s) and s != "None" and not s.startswith("SMILES:") and not s.startswith("NATURAL:")

    def get_combined_structural_token(self, cas: str, smiles: str) -> np.ndarray:
        """Main entry point. Returns a 151-dimensional float32 array."""
        return self.get_embedding(smiles, cas=cas)

    def _extract_physical_features(self, record: dict[str, Any]) -> np.ndarray:
        """Return a 2-D normalized [MW, LogP] vector."""
        mw = float(record.get("molecular_weight", 150.0))
        logp_raw = record.get("logp")
        logp = float(logp_raw) if logp_raw is not None else self.global_median_logp
        return np.array([mw / 300.0, logp / 5.0], dtype=np.float32)

    def _extract_functional_features(self, record: dict[str, Any]) -> np.ndarray:
        """Return a 4-D binary flag vector for ester/alcohol/ketone/aldehyde."""
        return _functional_group_flags(record.get("unifac_groups"))

    def get_embedding(self, smiles: str, cas: str | None = None) -> np.ndarray:
        """Main entry point. Returns a float32 embedding (dim = self.embedding_dim)."""
        normalized_cas = self._normalized_cas(cas)
        is_natural = bool(cas and str(cas).startswith("NATURAL:"))

        if normalized_cas:
            try:
                record = self.registry.get(normalized_cas)
                if record:
                    if self.use_real_pom or self.use_disjoint_pom:
                        # POM/disjoint arms: never reuse the registry's 138-dim Morgan block.
                        structural = self._compute_structural_embedding(smiles, cas)
                    elif self.use_morgan_2048_rp or self.use_morgan_2048_direct:
                        # 2048-bit arms: always fresh-compute the wide fingerprint;
                        # the registry cache only holds the 138-D block.
                        structural = self._compute_structural_embedding(smiles, cas)
                    else:
                        structural = np.asarray(record.get("openpom_embedding") or [0.0] * STRUCTURAL_DIM, dtype=np.float32)
                        if structural.shape[0] != STRUCTURAL_DIM or (not is_natural and np.count_nonzero(structural) == 0):
                            structural = self._compute_structural_embedding(smiles, cas)
                    shape = np.asarray(record.get("shape_3d_features") or [0.0] * SHAPE_DIM, dtype=np.float32)
                    if shape.shape[0] != SHAPE_DIM or (not is_natural and np.count_nonzero(shape) == 0):
                        shape = self._compute_shape_features(smiles, cas)
                    physical = self._extract_physical_features(record)
                    functional = self._extract_functional_features(record)
                    combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32)
                    combined = self._maybe_project(combined)
                    norm = np.linalg.norm(combined)
                    return combined / norm if norm > 0 else combined
                if is_natural:
                    structural = self._compute_structural_embedding(smiles, cas)
                    shape = self._compute_shape_features(smiles, cas)
                    physical = np.array([0.5, self.global_median_logp / 5.0], dtype=np.float32)
                    functional = _functional_group_flags(None)
                    combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32)
                    combined = self._maybe_project(combined)
                    norm = np.linalg.norm(combined)
                    return combined / norm if norm > 0 else combined
            except Exception as exc:
                logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc)

        # Fallback: compute structural/shape features on the fly.
        structural = self._compute_structural_embedding(smiles, cas)
        shape = self._compute_shape_features(smiles, cas)
        physical = np.array([0.5, self.global_median_logp / 5.0], dtype=np.float32)
        functional = _functional_group_flags(None)
        combined = np.concatenate([structural, shape, physical, functional]).astype(np.float32)
        combined = self._maybe_project(combined)
        norm = np.linalg.norm(combined)
        return combined / norm if norm > 0 else combined

    def _maybe_project(self, combined: np.ndarray) -> np.ndarray:
        """morgan_2048_rp arm: fixed seeded random projection 2071-D -> 269-D.

        The projection is generated from a fixed seed on first use (no learned
        parameters, no checkpoint dependency), so every process computes the
        identical map. Output width matches the OpenPOM arm's combined token,
        giving the two trained arms an identical input-projection parameter
        count (capacity parity for the representation comparison). The
        morgan_2048_direct arm skips this projection: its 2071-D token feeds a
        LEARNED input projection directly (round-10 panel response: the RP
        control constrains the usable subspace; the direct arm tests whether
        learning from the raw 2048 bits changes the ordering)."""
        if not self.use_morgan_2048_rp:
            return combined
        if self._rp_matrix is None:
            rng = np.random.default_rng(MORGAN2048_RP_SEED)
            self._rp_matrix = (
                rng.standard_normal((combined.shape[0], COMBINED_DIM_REAL_POM)).astype(np.float32)
                / np.sqrt(float(combined.shape[0]))
            )
        return combined @ self._rp_matrix

    def _compute_structural_embedding(self, smiles: str, cas: str | None = None) -> np.ndarray:
        normalized_cas = self._normalized_cas(cas)

        # 2048-bit arms: always the wide fingerprint block (never the cached
        # 138-D registry block, which is a different representation).
        if self.use_morgan_2048_rp or self.use_morgan_2048_direct:
            return self._morgan_structural_wide(smiles)

        # Genuine OpenPOM arm / disjoint-encoder arm: authoritative 256-dim
        # embedding for known CAS (both assets share the CAS-indexed layout and
        # are loaded into the same slots; exactly one arm is active at a time).
        if self.use_real_pom or self.use_disjoint_pom:
            if normalized_cas and normalized_cas in self._pom_index and self._pom_matrix is not None:
                return self._pom_matrix[self._pom_index[normalized_cas]]
            # CAS absent from the asset: real Morgan block, zero-padded
            # to the 256-dim structural slot so the combined dim stays constant.
            morgan = self._morgan_structural(smiles)
            padded = np.zeros(self.structural_dim, dtype=np.float32)
            padded[: morgan.shape[0]] = morgan
            return padded

        # Morgan arm: live SQLite cache, natural-oil blend, then Morgan compute.
        if normalized_cas:
            try:
                record = self.registry.get(normalized_cas)
                if record:
                    openpom_embedding = record.get("openpom_embedding")
                    if openpom_embedding and len(openpom_embedding) > 0:
                        return np.asarray(openpom_embedding, dtype=np.float32)
            except Exception as exc:
                logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc)

        if normalized_cas and ((cas or "").startswith("NATURAL:") or normalized_cas.startswith(("8", "6", "9"))):
            try:
                all_records = self.registry.all_records()
                openpom, _ = resolve_natural_oil_vectors(normalized_cas, all_records)
                if any(openpom):
                    return np.asarray(openpom, dtype=np.float32)
            except Exception as exc:
                logger.debug("Natural oil resolution failed for %s: %s", normalized_cas, exc)

        return self._morgan_structural(smiles)

    def _morgan_structural_wide(self, smiles: str) -> np.ndarray:
        """Deterministic 2048-bit radius-2 Morgan + 10 physicochemical descriptors (2058-D).

        The standard-width fingerprint control: identical construction to the
        preregistered 128-bit baseline, only fpSize differs."""
        if not self._is_valid_smiles(smiles):
            return np.zeros(MORGAN2048_STRUCTURAL_DIM, dtype=np.float32)
        mol = Chem.MolFromSmiles(smiles)
        if not mol:
            return np.zeros(MORGAN2048_STRUCTURAL_DIM, dtype=np.float32)
        try:
            fpgen = AllChem.GetMorganGenerator(radius=2, fpSize=2048)
            fp_array = np.array(fpgen.GetFingerprint(mol), dtype=np.float32)
        except AttributeError:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048)
            fp_array = np.array(fp, dtype=np.float32)
        physchem = np.array([
            Descriptors.MolWt(mol), Descriptors.ExactMolWt(mol), Descriptors.MolLogP(mol),
            Descriptors.HeavyAtomCount(mol), Descriptors.NumHAcceptors(mol),
            Descriptors.NumHDonors(mol), Descriptors.NumRotatableBonds(mol),
            Descriptors.FractionCSP3(mol), Descriptors.TPSA(mol), Descriptors.LabuteASA(mol),
        ], dtype=np.float32)
        return np.concatenate([fp_array, physchem]).astype(np.float32)

    def _morgan_structural(self, smiles: str) -> np.ndarray:
        """Deterministic 128-bit Morgan + 10 physicochemical descriptors (138-D).

        Used as the Morgan arm's structural block and as the fallback for CAS
        absent from the genuine-POM asset."""
        if not self._is_valid_smiles(smiles):
            return np.zeros(STRUCTURAL_DIM, dtype=np.float32)
        mol = Chem.MolFromSmiles(smiles)
        if not mol:
            return np.zeros(STRUCTURAL_DIM, dtype=np.float32)
        try:
            fpgen = AllChem.GetMorganGenerator(radius=2, fpSize=128)
            fp_array = np.array(fpgen.GetFingerprint(mol), dtype=np.float32)
        except AttributeError:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=128)
            fp_array = np.array(fp, dtype=np.float32)
        physchem = np.array([
            Descriptors.MolWt(mol), Descriptors.ExactMolWt(mol), Descriptors.MolLogP(mol),
            Descriptors.HeavyAtomCount(mol), Descriptors.NumHAcceptors(mol),
            Descriptors.NumHDonors(mol), Descriptors.NumRotatableBonds(mol),
            Descriptors.FractionCSP3(mol), Descriptors.TPSA(mol), Descriptors.LabuteASA(mol),
        ], dtype=np.float32)
        return np.concatenate([fp_array, physchem]).astype(np.float32)

    def _compute_shape_features(self, smiles: str, cas: str | None = None) -> np.ndarray:
        """Fetch cached 3D shape features or compute them on the fly and write through."""
        normalized_cas = self._normalized_cas(cas)
        zero_shape = np.zeros(self.shape_dim, dtype=np.float32)

        # 1. Query the live SQLite cache.
        if normalized_cas:
            try:
                record = self.registry.get(normalized_cas)
                if record:
                    shape_3d = record.get("shape_3d_features")
                    if shape_3d and len(shape_3d) > 0 and any(v != 0.0 for v in shape_3d):
                        return np.asarray(shape_3d, dtype=np.float32)
            except Exception as exc:
                logger.debug("Registry lookup failed for %s: %s", normalized_cas, exc)

        # Natural oils may be resolvable from a constituent blend.
        if normalized_cas and (normalized_cas.startswith("8") or normalized_cas.startswith("6") or normalized_cas.startswith("9")):
            try:
                all_records = self.registry.all_records()
                _, shape = resolve_natural_oil_vectors(normalized_cas, all_records)
                if any(shape):
                    return np.asarray(shape, dtype=np.float32)
            except Exception as exc:
                logger.debug("Natural oil shape resolution failed for %s: %s", normalized_cas, exc)

        # 2. String guardrail: invalid/empty SMILES short-circuit to zero fallback.
        if not self._is_valid_smiles(smiles):
            return zero_shape

        # 3. Just-In-Time compute 3D features exactly once.
        try:
            features = extract_3d_shape_features(smiles)
        except Exception as exc:
            logger.debug("3D shape feature extraction failed for %s: %s", smiles, exc)
            return zero_shape

        shape_3d_vec = list(features)
        if any(v != 0.0 for v in shape_3d_vec):
            # 4. Write through to the persistent SQLite cache.
            if normalized_cas:
                try:
                    record = self.registry.get(normalized_cas) or {"cas": normalized_cas}
                    record["shape_3d_features"] = shape_3d_vec
                    self.registry.add(record)
                except Exception as exc:
                    logger.debug("Failed to persist 3D features for %s: %s", normalized_cas, exc)

        return np.asarray(shape_3d_vec, dtype=np.float32)

    def _compute_real_openpom(self, smiles: str) -> np.ndarray:
        raise NotImplementedError("Real OpenPOM model not yet loaded")


# Module-level engine is lazily instantiated so that import-time side effects
# (e.g., opening the SQLite registry) do not run until the first actual use.
_DEFAULT_ENGINE: OlfactoryEmbeddingEngine | None = None


def _get_default_engine() -> OlfactoryEmbeddingEngine:
    global _DEFAULT_ENGINE
    if _DEFAULT_ENGINE is None:
        _DEFAULT_ENGINE = OlfactoryEmbeddingEngine(use_fallback=True)
    return _DEFAULT_ENGINE


def get_openpom_embedding(smiles: str, dim: int = COMBINED_DIM, cas: str | None = None) -> np.ndarray:
    """Convenience accessor used by the PIMT pipeline."""
    engine = _get_default_engine()
    embedding = engine.get_embedding(smiles, cas=cas)
    if embedding.shape[0] != dim:
        raise ValueError(f"Embedding dimension mismatch: {embedding.shape[0]} != {dim}")
    return embedding