src/pino/registry.py

from __future__ import annotations

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

import numpy as np

from . import unifac_mapping
from .models import IngredientResolutionError
from .thermo.geometry import extract_3d_shape_features

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


class AromaRegistry:
    """
    SQLite-backed registry of pre-computed structural and thermodynamic data for
    common fragrance ingredients. Used as a fast lookup before falling back to
    live ugropy/RDKit fragmentation.

    Schema:
        cas TEXT PRIMARY KEY,
        name TEXT,
        smiles TEXT,
        molecular_weight REAL,
        boiling_point_k REAL,
        vapor_pressure_pa REAL,
        odor_threshold_ug_m3 REAL,
        logp REAL,
        unifac_groups TEXT  -- JSON dict of string subgroup_id -> count
    """

    DEFAULT_PATH = Path(__file__).with_suffix(".db")

    def __init__(self, path: Path | str | None = None) -> None:
        self.path = Path(path) if path else Path(os.environ.get("PINO_REGISTRY_PATH", self.DEFAULT_PATH))
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(self.path)
        self._conn.row_factory = sqlite3.Row
        self._create_tables()

    def _create_tables(self) -> None:
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS aroma_chemicals (
                cas TEXT PRIMARY KEY,
                name TEXT,
                smiles TEXT,
                molecular_weight REAL,
                boiling_point_k REAL,
                vapor_pressure_pa REAL,
                odor_threshold_ug_m3 REAL,
                logp REAL,
                odor_description TEXT,
                unifac_groups TEXT,
                openpom_embedding TEXT,
                shape_3d_features TEXT
            )
            """
        )
        # Migrate older registries that were created without optional columns.
        for column in ("odor_description", "openpom_embedding", "shape_3d_features"):
            try:
                self._conn.execute(f"ALTER TABLE aroma_chemicals ADD COLUMN {column} TEXT")
                self._conn.commit()
            except sqlite3.OperationalError:
                pass  # Column already exists
        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_name ON aroma_chemicals(name)"
        )
        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_smiles ON aroma_chemicals(smiles)"
        )
        self._conn.commit()

    def register(
        self,
        cas: str,
        name: str,
        smiles: str,
        molecular_weight: float,
        boiling_point_k: float | None = None,
        vapor_pressure_pa: float | None = None,
        logp: float | None = None,
        unifac_groups_json: str = "{}",
        source: str = "manual",
    ) -> None:
        """Insert a fully-built record directly (used by the expander)."""
        self._conn.execute(
            """
            INSERT OR REPLACE INTO aroma_chemicals
            (cas, name, smiles, molecular_weight, boiling_point_k, vapor_pressure_pa,
             odor_threshold_ug_m3, logp, unifac_groups)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                cas,
                name,
                smiles,
                molecular_weight,
                boiling_point_k,
                vapor_pressure_pa,
                None,
                logp,
                json.dumps(unifac_groups_json, sort_keys=True),
            ),
        )
        self._conn.commit()

    def close(self) -> None:
        self._conn.close()

    def __enter__(self) -> AromaRegistry:
        return self

    def __exit__(self, *args) -> None:
        self.close()

    def _normalise_key(self, key: str) -> str:
        return str(key).strip().lower()

    def get(self, identifier: str) -> dict[str, Any] | None:
        """Look up a molecule by CAS, name, or SMILES."""
        key = self._normalise_key(identifier)
        for column in ("cas", "name", "smiles"):
            row = self._conn.execute(
                f"""SELECT * FROM aroma_chemicals
                    WHERE LOWER({column}) = ?
                    ORDER BY
                        CASE WHEN cas LIKE 'SMILES:%' THEN 1 ELSE 0 END,
                        CASE WHEN vapor_pressure_pa IS NULL THEN 1 ELSE 0 END,
                        CASE WHEN boiling_point_k IS NULL THEN 1 ELSE 0 END,
                        cas
                    LIMIT 1""",
                (key,),
            ).fetchone()
            if row:
                record = dict(row)
                record["unifac_groups"] = json.loads(record.get("unifac_groups") or "{}")
                record["openpom_embedding"] = json.loads(record.get("openpom_embedding") or "[]")
                record["shape_3d_features"] = json.loads(record.get("shape_3d_features") or "[]")
                return record
        return None

    def __contains__(self, identifier: str) -> bool:
        return self.get(identifier) is not None

    def add(self, record: dict[str, Any]) -> None:
        """Insert or replace a registry record keyed by CAS."""
        smiles = record.get("smiles", "")
        cas = record.get("cas", "")
        openpom = record.get("openpom_embedding")
        shape_3d = record.get("shape_3d_features")

        if openpom is None and smiles and not smiles.startswith("NATURAL:"):
            from .embeddings import OlfactoryEmbeddingEngine
            engine = OlfactoryEmbeddingEngine(use_fallback=True)
            openpom = engine._compute_structural_embedding(smiles, cas=cas)
        if shape_3d is None and smiles and not smiles.startswith("NATURAL:"):
            shape_3d = extract_3d_shape_features(smiles)

        if isinstance(openpom, (list, tuple, np.ndarray)):
            openpom = json.dumps(np.asarray(openpom).tolist())
        if isinstance(shape_3d, (list, tuple, np.ndarray)):
            shape_3d = json.dumps(np.asarray(shape_3d).tolist())

        self._conn.execute(
            """
            INSERT OR REPLACE INTO aroma_chemicals
            (cas, name, smiles, molecular_weight, boiling_point_k, vapor_pressure_pa,
             odor_threshold_ug_m3, logp, odor_description, unifac_groups, openpom_embedding, shape_3d_features)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                record["cas"],
                record.get("name"),
                record.get("smiles"),
                record.get("molecular_weight"),
                record.get("boiling_point_k"),
                record.get("vapor_pressure_pa"),
                record.get("odor_threshold_ug_m3"),
                record.get("logp"),
                record.get("odor_description", ""),
                json.dumps(record.get("unifac_groups", {}), sort_keys=True),
                openpom,
                shape_3d,
            ),
        )
        self._conn.commit()

    @staticmethod
    def validate_smiles(smiles: str) -> dict[str, Any] | None:
        """
        Validate a SMILES string locally: RDKit parse, MW guard, ugropy
        Dortmund-UNIFAC fragmentation, and thermo parameter compatibility.
        Returns a registry-ready record dict or None if validation fails.
        """
        from rdkit import Chem
        from rdkit.Chem import Descriptors

        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        try:
            mw = float(Descriptors.MolWt(mol))
        except Exception:
            return None

        try:
            logp = float(Descriptors.MolLogP(mol))
        except Exception:
            logp = None

        try:
            groups = AromaRegistry.fragment_groups(smiles, "smiles")
        except Exception:
            # Some molecules (e.g. coumarin) parse in RDKit but cannot be
            # fragmented by ugropy. We still keep the record so that VLE can
            # fall back to ideal gamma=1.0 rather than rejecting the material.
            groups = {}

        return {
            "smiles": Chem.MolToSmiles(mol, canonical=True),
            "molecular_weight": mw,
            "logp": logp,
            "unifac_groups": groups,
            "name": smiles,
        }

    @staticmethod
    def build_record_from_smiles(
        smiles: str,
        name: str | None = None,
        vapor_pressure_pa: float | None = None,
        boiling_point_k: float | None = None,
        odor_threshold_ug_m3: float | None = None,
    ) -> dict[str, Any]:
        """
        Build a registry record entirely offline from a SMILES string.
        No PubChem round-trip; CAS is generated as a synthetic SMILES key if no
        human-readable name is supplied. RDKit provides MW and LogP, ugropy
        provides Dortmund-UNIFAC groups, and thermo validates that every
        subgroup has interaction parameters.
        """
        from rdkit import Chem
        from rdkit.Chem import Descriptors

        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            raise IngredientResolutionError(f"Invalid SMILES: {smiles}", smiles=smiles)
        canonical = Chem.MolToSmiles(mol, canonical=True)
        mw = float(Descriptors.MolWt(mol))
        try:
            logp = float(Descriptors.MolLogP(mol))
        except Exception:
            logp = None
        groups = AromaRegistry.fragment_groups(canonical, "smiles")
        if not groups:
            # Allow ideal-solution fallback for materials that cannot be fragmented
            # (e.g., coumarin). The VLE model will use gamma=1.0.
            groups = {}
        return {
            "cas": f"SMILES:{canonical}",
            "name": name or canonical,
            "smiles": canonical,
            "molecular_weight": mw,
            "vapor_pressure_pa": vapor_pressure_pa,
            "boiling_point_k": boiling_point_k,
            "odor_threshold_ug_m3": odor_threshold_ug_m3,
            "logp": logp,
            "unifac_groups": groups,
        }

    @staticmethod
    def resolve_pubchem(identifier: str, identifier_type: str = "cas") -> dict[str, Any]:
        """Resolve a molecule to canonical name, SMILES, and MW via PubChem."""
        try:
            import pubchempy as pcp

            if identifier_type == "cas":
                compounds = pcp.get_compounds(identifier, "name")
            else:
                compounds = pcp.get_compounds(identifier, identifier_type)
        except Exception as exc:
            raise IngredientResolutionError(
                f"PubChem lookup failed for {identifier}: {exc}", cas=identifier
            ) from exc

        if not compounds:
            raise IngredientResolutionError(
                f"PubChem returned no compound for {identifier}", cas=identifier
            )

        comp = compounds[0]
        return {
            "cas": identifier,
            "name": comp.synonyms[0] if comp.synonyms else identifier,
            "smiles": comp.canonical_smiles or comp.smiles,
            "molecular_weight": float(comp.molecular_weight),
        }

    @staticmethod
    def _normalise_ugropy_name(name: str) -> str:
        """Map ugropy subgroup names onto the thermo DDB UNIFAC namespace."""
        aliases = {
            "HCO": "CHO",  # aldehyde: same group, different label
            "OH (P)": "OH(P)",
            "OH (S)": "OH(S)",
            "OH (T)": "OH(T)",
            "CH=O": "CHO",
        }
        return aliases.get(name, name).replace(" ", "")

    @staticmethod
    def fragment_groups(identifier: str, identifier_type: str = "name") -> dict[str, int]:
        """Run ugropy and return thermo-compatible integer subgroup IDs."""
        from thermo.unifac import DOUFSG

        import ugropy

        name_to_id = {str(v.group): k for k, v in DOUFSG.items()}
        try:
            groups_obj = ugropy.Groups(identifier, identifier_type=identifier_type)
            raw_groups = groups_obj.dortmund.subgroups
        except Exception as exc:
            raise IngredientResolutionError(
                f"ugropy fragmentation failed for {identifier}: {exc}",
                cas=identifier if identifier_type == "cas" else None,
            ) from exc

        if not raw_groups:
            raise IngredientResolutionError(
                f"ugropy returned no Dortmund groups for {identifier}",
                cas=identifier if identifier_type == "cas" else None,
            )

        result: dict[str, int] = {}
        for name, count in raw_groups.items():
            subgroup_id = unifac_mapping.map_ugropy_to_thermo_id(name, name_to_id)
            if subgroup_id is None:
                raise IngredientResolutionError(
                    f"Dortmund subgroup '{name}' from {identifier} not in thermo parameters"
                )
            result[str(subgroup_id)] = int(count)
        return result

    @staticmethod
    def build_record(
        identifier: str,
        identifier_type: str = "cas",
        vapor_pressure_pa: float | None = None,
        boiling_point_k: float | None = None,
        odor_threshold_ug_m3: float | None = None,
        logp: float | None = None,
    ) -> dict[str, Any]:
        """Build a registry record by resolving PubChem and fragmenting groups."""
        from rdkit import Chem
        from rdkit.Chem import Descriptors

        base = AromaRegistry.resolve_pubchem(identifier, identifier_type)
        exc: Exception | None = None
        for id_for_ugropy, id_type in [
            (identifier, identifier_type),
            (base["name"], "name"),
            (base["smiles"], "smiles"),
        ]:
            try:
                groups = AromaRegistry.fragment_groups(id_for_ugropy, id_type)
                break
            except Exception as e:
                exc = e
                groups = {}
        else:
            raise exc or IngredientResolutionError(
                f"Could not fragment {identifier} by CAS, name, or SMILES"
            )

        # Compute LogP from RDKit if not provided.
        if logp is None:
            try:
                mol = Chem.MolFromSmiles(base["smiles"])
                logp = float(Descriptors.MolLogP(mol)) if mol else None
            except Exception:
                logp = None

        base["unifac_groups"] = groups
        base["vapor_pressure_pa"] = vapor_pressure_pa
        base["boiling_point_k"] = boiling_point_k
        base["odor_threshold_ug_m3"] = odor_threshold_ug_m3
        base["logp"] = logp
        return base

    def all_records(self) -> dict[str, dict[str, Any]]:
        """Return all registry rows keyed by CAS for bulk lookups."""
        rows = self._conn.execute("SELECT * FROM aroma_chemicals").fetchall()
        return {
            row["cas"]: {
                **dict(row),
                "unifac_groups": json.loads(row["unifac_groups"] or "{}"),
                "shape_3d_features": json.loads(row["shape_3d_features"] or "[]"),
                "openpom_embedding": json.loads(row["openpom_embedding"] or "[]"),
            }
            for row in rows
        }

    def backfill_3d_shape_features(self) -> None:
        """Compute and store 3D shape features for all registry rows lacking them."""
        rows = self._conn.execute(
            "SELECT cas, smiles FROM aroma_chemicals WHERE shape_3d_features IS NULL OR shape_3d_features = ?",
            (json.dumps([]),),
        ).fetchall()
        logger.info("Backfilling 3D shape features for %d registry entries", len(rows))
        for cas, smiles in rows:
            if not smiles or smiles.startswith("NATURAL:"):
                continue
            features = extract_3d_shape_features(smiles)
            self._conn.execute(
                "UPDATE aroma_chemicals SET shape_3d_features = ? WHERE cas = ?",
                (json.dumps(features), cas),
            )
        self._conn.commit()
        logger.info("3D shape feature backfill complete")

    def backfill_openpom_embeddings(self) -> None:
        """Compute and store 138-D OpenPOM embeddings for all single-molecule rows."""
        from .embeddings import OlfactoryEmbeddingEngine

        rows = self._conn.execute(
            "SELECT cas, smiles FROM aroma_chemicals WHERE openpom_embedding IS NULL OR openpom_embedding = ?",
            (json.dumps([]),),
        ).fetchall()
        logger.info("Backfilling OpenPOM embeddings for %d registry entries", len(rows))
        engine = OlfactoryEmbeddingEngine(use_fallback=True)
        for cas, smiles in rows:
            if not smiles or smiles.startswith("NATURAL:"):
                continue
            structural = engine._compute_structural_embedding(smiles, cas=cas)
            self._conn.execute(
                "UPDATE aroma_chemicals SET openpom_embedding = ? WHERE cas = ?",
                (json.dumps(structural.tolist()), cas),
            )
        self._conn.commit()
        logger.info("OpenPOM embedding backfill complete")

    def backfill_natural_oil_vectors(self) -> None:
        """Two-pass: build weighted OpenPOM + shape vectors for mapped natural oils."""
        from .thermo.naturals import resolve_natural_oil_vectors

        rows = self._conn.execute(
            "SELECT cas FROM aroma_chemicals WHERE (cas LIKE '8000-%' OR cas LIKE '8007-%' OR cas LIKE '8014-%' OR cas LIKE '8016-%' OR cas LIKE '8022-%' OR cas LIKE '8023-%' OR cas LIKE '8024-%' OR cas LIKE '8031-%' OR cas LIKE '8046-%' OR cas LIKE '9000-%' OR cas LIKE '68606-%' OR cas LIKE '68855-%' OR cas LIKE '72968-%' OR cas LIKE '89958-%' OR cas LIKE '90045-%') AND (openpom_embedding IS NULL OR shape_3d_features IS NULL)"
        ).fetchall()
        logger.info("Backfilling natural oil vectors for %d entries", len(rows))
        cache = self.all_records()
        for (cas,) in rows:
            openpom, shape = resolve_natural_oil_vectors(cas, cache)
            self._conn.execute(
                "UPDATE aroma_chemicals SET openpom_embedding = ?, shape_3d_features = ? WHERE cas = ?",
                (json.dumps(openpom), json.dumps(shape), cas),
            )
        self._conn.commit()
        logger.info("Natural oil vector backfill complete")

    def populate(
        self,
        entries: list[dict[str, Any]],
        *,
        skip_failures: bool = True,
    ) -> list[dict[str, Any]]:
        """
        Populate the registry from a list of entries. Each entry is a dict with
        at least "cas" and optionally "vapor_pressure_pa", "boiling_point_k",
        "odor_threshold_ug_m3", "logp".
        """
        failed: list[dict[str, Any]] = []
        for entry in entries:
            cas = entry["cas"]
            try:
                record = self.build_record(
                    cas,
                    "cas",
                    vapor_pressure_pa=entry.get("vapor_pressure_pa"),
                    boiling_point_k=entry.get("boiling_point_k"),
                    odor_threshold_ug_m3=entry.get("odor_threshold_ug_m3"),
                    logp=entry.get("logp"),
                )
                self.add(record)
                logger.info("Added registry entry for CAS %s (%s)", cas, record.get("name"))
            except Exception as exc:
                logger.warning("Failed to build registry entry for CAS %s: %s", cas, exc)
                failed.append(entry)
                if not skip_failures:
                    raise
        return failed