src/pino/palette_expander.py

from __future__ import annotations

import argparse
import json
import logging
import sqlite3
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

import pandas as pd
import pyrfume
from rdkit import Chem
from rdkit.Chem import Descriptors

from pino.registry import AromaRegistry

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

PROHIBITED_CAS = {
    "81-15-2": "Musk xylene",
    "83-66-9": "Musk ambrette",
}


def _load_pyrfume_molecules(archives: list[str]) -> pd.DataFrame:
    """Load and deduplicate molecule metadata from pyrfume archives."""
    frames = []
    for archive in archives:
        try:
            df = pyrfume.load_data(f"{archive}/molecules.csv")
            if df is None or df.empty:
                continue
            df = df.reset_index().rename(columns={"index": "CID"})
            df["archive"] = archive
            frames.append(df)
            logger.info("Loaded %s: %d molecules", archive, len(df))
        except Exception as exc:
            logger.warning("Could not load %s: %s", archive, exc)
    if not frames:
        raise RuntimeError("No pyrfume archives could be loaded")
    combined = pd.concat(frames, ignore_index=True)
    combined = combined.drop_duplicates(subset=["IsomericSMILES"], keep="first")
    return combined


def _validate_molecule(smiles: str, name: str | None = None) -> dict[str, Any] | None:
    """Offline SMILES validation through RDKit + ugropy + thermo."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    try:
        mw = Descriptors.MolWt(mol)
    except Exception:
        return None
    if mw < 50 or mw > 600:
        return None
    try:
        record = AromaRegistry.build_record_from_smiles(smiles, name=name)
    except Exception as exc:
        logger.debug("Validation failed for %s: %s", smiles, exc)
        return None
    return record


def _harvest_worker(args: tuple[int, pd.Series]) -> dict[str, Any] | None:
    """Worker for a single molecule: validate and build a registry record."""
    idx, row = args
    smiles = row.get("IsomericSMILES")
    if not isinstance(smiles, str) or not smiles:
        return None
    name = row.get("name") or row.get("IUPACName")
    record = _validate_molecule(smiles, name=name)
    if record is None:
        return None

    cas = record.get("cas", "")
    if cas in PROHIBITED_CAS:
        logger.info("Skipping prohibited compound %s (%s)", cas, PROHIBITED_CAS[cas])
        return None

    record["source"] = f"pyrfume:{row.get('archive', 'unknown')}"
    return record


def expand_registry(
    registry_path: str | Path,
    target: int = 300,
    max_workers: int = 8,
    batch_size: int = 50,
) -> dict[str, int]:
    """
    Harvest aroma molecules from pyrfume, validate them, and insert into the
    SQLite registry until `target` new validated records are reached.

    This routine is checkpoint-resilient: existing records are skipped on
    restart, and successful inserts are committed in batches.
    """
    registry = AromaRegistry(Path(registry_path))
    existing = {
        r["smiles"]
        for r in registry._conn.execute("SELECT smiles FROM aroma_chemicals").fetchall()
    }
    logger.info("Registry currently has %d records; skipping known SMILES", len(existing))

    archives = ["goodscents", "leffingwell", "arctander_1960", "fragrancedb"]
    molecules = _load_pyrfume_molecules(archives)
    # Drop already-known SMILES before doing any expensive validation.
    molecules = molecules[~molecules["IsomericSMILES"].isin(existing)]
    logger.info("Candidate molecules after deduplication and skip: %d", len(molecules))

    inserted = 0
    rejected = 0
    skipped = 0
    pending: list[dict[str, Any]] = []

    def _commit_batch(records: list[dict[str, Any]]) -> int:
        """Insert a batch of validated records and return the count inserted."""
        count = 0
        for record in records:
            try:
                registry.register(
                    cas=record["cas"],
                    name=record["name"],
                    smiles=record["smiles"],
                    molecular_weight=record["molecular_weight"],
                    boiling_point_k=record.get("boiling_point_k"),
                    vapor_pressure_pa=record.get("vapor_pressure_pa"),
                    logp=record.get("logp"),
                    unifac_groups_json=record.get("unifac_groups", {}),
                )
                count += 1
            except Exception as exc:
                logger.warning("DB insert failed for %s: %s", record.get("cas"), exc)
        return count

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(_harvest_worker, (i, row)): (i, row)
            for i, row in molecules.iterrows()
        }
        for future in as_completed(futures):
            record = future.result()
            if record is None:
                rejected += 1
                continue
            if record["smiles"] in existing:
                skipped += 1
                continue
            pending.append(record)
            existing.add(record["smiles"])

            if len(pending) >= batch_size:
                inserted += _commit_batch(pending)
                pending = []
                logger.info(
                    "Checkpoint: inserted=%d rejected=%d skipped=%d | registry total ~%d",
                    inserted,
                    rejected,
                    skipped,
                    len(existing),
                )
            if inserted >= target:
                break

    # Final flush.
    if pending:
        inserted += _commit_batch(pending)

    registry.close()
    return {"inserted": inserted, "rejected": rejected, "skipped": skipped}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Expand PINO aroma-chemical registry from pyrfume")
    parser.add_argument("--registry", default="src/pino/registry.db", help="Path to SQLite registry")
    parser.add_argument("--target", type=int, default=300, help="Target number of new validated records")
    parser.add_argument("--workers", type=int, default=8, help="Parallel validation workers")
    parser.add_argument("--batch-size", type=int, default=50, help="SQLite commit batch size")
    parser.add_argument("--log-level", default="INFO", help="Logging level")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    result = expand_registry(args.registry, args.target, args.workers, args.batch_size)
    print(result)