src/pino/synthetic.py

from __future__ import annotations

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

import numpy as np

from pino.ifra import IFRA_RESTRICTIONS
from pino.registry import AromaRegistry
from pino.verifier import FragrancePipelineVerifier

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


def _light_ifra_check(formula: list[dict[str, Any]]) -> dict[str, Any]:
    """Fast pre-check against IFRA restrictions using raw dict input."""
    violations = []
    for f in formula:
        cas = f["cas"]
        pct = f["weight_fraction"] * 100.0
        if cas in IFRA_RESTRICTIONS:
            limit = IFRA_RESTRICTIONS[cas]["category_4_pct"]
            if (limit == 0.0 and pct > 0.0) or (limit > 0.0 and pct > limit):
                violations.append({"cas": cas, "used_pct": pct, "limit_pct": limit})
    return {"passed": not violations, "violations": violations}


class SyntheticFormulationGenerator:
    """
    High-throughput generator of physically plausible fragrance formulations.

    Samples ingredient palettes from the validated aroma registry, applies
    realistic concentration guardrails, filters through IFRA and the VLE
    verifier, and returns a clean ML-ready corpus.
    """

    def __init__(
        self,
        registry_path: str | Path = "src/pino/registry.db",
        seed: int | None = None,
        min_ingredients: int = 4,
        max_ingredients: int = 10,
        base_fraction_range: tuple[float, float] = (0.50, 0.90),
        accord_mass_fraction: float = 0.30,
        max_attempts: int = 100,
    ) -> None:
        self.registry_path = Path(registry_path)
        self.seed = seed
        self.rng = random.Random(seed)
        self.min_ingredients = min_ingredients
        self.max_ingredients = max_ingredients
        self.base_fraction_range = base_fraction_range
        self.accord_mass_fraction = accord_mass_fraction
        self.max_attempts = max_attempts
        self._registry = AromaRegistry(str(self.registry_path))
        self._verifier = FragrancePipelineVerifier()

    def _load_aroma_chemicals(self) -> list[dict[str, Any]]:
        rows = self._registry._conn.execute(
            "SELECT cas, name, molecular_weight, vapor_pressure_pa, boiling_point_k, logp FROM aroma_chemicals"
        ).fetchall()
        return [
            {
                "cas": row[0],
                "name": row[1],
                "molecular_weight": row[2],
                "vapor_pressure_pa": row[3],
                "boiling_point_k": row[4],
                "logp": row[5],
            }
            for row in rows
        ]

    def _pick_base(self, chemicals: list[dict[str, Any]]) -> dict[str, Any]:
        # Prefer ethanol as the canonical volatile base if available; otherwise fallback.
        for chem in chemicals:
            if chem["cas"] == "64-17-5":
                return chem
        return self.rng.choice(chemicals)

    def _sample_palette(self, chemicals: list[dict[str, Any]], base: dict[str, Any]) -> list[dict[str, Any]]:
        n = self.rng.randint(self.min_ingredients, self.max_ingredients)
        # Ensure base is included, then fill with diverse aroma chemicals.
        candidates = [c for c in chemicals if c["cas"] != base["cas"]]
        if len(candidates) < n - 1:
            n = len(candidates) + 1
        selected = self.rng.sample(candidates, n - 1)
        return [base] + selected

    def _distribute_weights(self, palette: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """
        Realistic fragrance pyramid:
          - Base solvent dominates (e.g., ethanol).
          - Remaining mass is distributed across the accord with a slight bias toward
            lower vapor pressure (heart/base) notes for stability.
        """
        base_cas = palette[0]["cas"]
        base_fraction = self.rng.uniform(*self.base_fraction_range)
        accord_mass = 1.0 - base_fraction

        # Use inverse volatility weighting for the accord to avoid all-top-note formulas.
        volatilities = np.array([max(c["vapor_pressure_pa"] or 1e-12, 1e-12) for c in palette[1:]])
        # Mix random and inverse volatility weights to keep diversity.
        random_weights = np.array([self.rng.random() for _ in palette[1:]])
        inverse_vol_weights = 1.0 / volatilities
        inverse_vol_weights /= inverse_vol_weights.sum()
        random_weights /= random_weights.sum()
        blend = 0.6 * random_weights + 0.4 * inverse_vol_weights
        blend /= blend.sum()
        accord_fractions = blend * accord_mass

        weights = {base_cas: base_fraction}
        for chem, frac in zip(palette[1:], accord_fractions):
            weights[chem["cas"]] = float(frac)

        formula = [{"cas": c["cas"], "weight_fraction": weights[c["cas"]]} for c in palette]
        return formula

    def _apply_guardrails(self, formula: list[dict[str, Any]]) -> bool:
        """
        Basic concentration guardrails before expensive VLE verification.
        """
        total = sum(f["weight_fraction"] for f in formula)
        if not np.isclose(total, 1.0, atol=1e-4):
            return False
        # No single accord ingredient dominates (>40% of non-base mass).
        base_cas = formula[0]["cas"]
        accord_total = sum(f["weight_fraction"] for f in formula if f["cas"] != base_cas)
        for f in formula:
            if f["cas"] != base_cas and f["weight_fraction"] / max(accord_total, 1e-12) > 0.40:
                return False
        return True

    def generate_one(self, duration_seconds: float = 28800, interval_seconds: float = 600) -> dict[str, Any] | None:
        """
        Generate and verify a single synthetic formulation. Returns the verified record
        or None if rejected/out-of-domain.
        """
        chemicals = self._load_aroma_chemicals()
        base = self._pick_base(chemicals)
        for attempt in range(self.max_attempts):
            palette = self._sample_palette(chemicals, base)
            formula = self._distribute_weights(palette)
            if not self._apply_guardrails(formula):
                continue
            # Lightweight IFRA pre-check before VLE.
            ifra_report = _light_ifra_check(formula)
            if not ifra_report["passed"]:
                continue
            try:
                result = self._verifier.run_sim(formula, duration_seconds=duration_seconds, interval_seconds=interval_seconds)
            except Exception as e:
                logger.debug("Verifier rejected formula: %s", e)
                continue
            if result["status"] in ("passed", "depleted"):
                return {
                    "formula": formula,
                    "status": result["status"],
                    "message": result["message"],
                    "depletion_rates": result["depletion_rates"],
                    "trajectory": result["trajectory"],
                }
        logger.warning("Failed to generate a verified formulation after %d attempts", self.max_attempts)
        return None

    def generate_corpus(
        self,
        n_samples: int = 100,
        duration_seconds: float = 28800,
        interval_seconds: float = 600,
        output_path: str | Path = "synthetic_corpus.json",
    ) -> dict[str, Any]:
        """
        Generate a corpus of n verified formulations and save as JSON.
        Returns a summary dict with pass/reject counts and output path.
        """
        output_path = Path(output_path)
        samples: list[dict[str, Any]] = []
        status_counts = {"passed": 0, "depleted": 0, "rejected": 0}
        for i in range(n_samples):
            record = self.generate_one(duration_seconds=duration_seconds, interval_seconds=interval_seconds)
            if record is None:
                status_counts["rejected"] += 1
                continue
            samples.append(record)
            status_counts[record["status"]] += 1
            if (i + 1) % 10 == 0:
                logger.info("Generated %d/%d verified formulations", len(samples), n_samples)
        corpus = {
            "metadata": {
                "n_samples": len(samples),
                "requested_samples": n_samples,
                "duration_seconds": duration_seconds,
                "interval_seconds": interval_seconds,
                "registry_path": str(self.registry_path),
                "seed": self.seed,
                "status_counts": status_counts,
            },
            "samples": samples,
        }
        output_path.write_text(json.dumps(corpus, indent=2))
        logger.info("Saved corpus with %d samples to %s", len(samples), output_path)
        return corpus


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate synthetic PINO training corpus")
    parser.add_argument("--n-samples", type=int, default=100, help="Number of verified formulations to generate")
    parser.add_argument("--duration", type=float, default=28800, help="Simulation duration in seconds")
    parser.add_argument("--interval", type=float, default=600, help="Output interval in seconds")
    parser.add_argument("--output", type=Path, default="synthetic_corpus.json", help="Output JSON path")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--registry", type=Path, default="src/pino/registry.db", help="Path to SQLite registry")
    parser.add_argument("--log-level", default="INFO", help="Logging level")
    args = parser.parse_args()

    logging.basicConfig(level=getattr(logging, args.log_level.upper()))
    gen = SyntheticFormulationGenerator(registry_path=args.registry, seed=args.seed)
    gen.generate_corpus(
        n_samples=args.n_samples,
        duration_seconds=args.duration,
        interval_seconds=args.interval,
        output_path=args.output,
    )