src/pino/draft_engine.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

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


class FormulaGenerator:
    """
    Genre-aware, tier-biased formulation sampler.

    Reads an expanded aroma-chemical registry, classifies compounds into
    performance tiers by molecular weight, and generates weight-fraction
    formulas that satisfy genre-specific blending rules.
    """

    def __init__(
        self,
        genre: str,
        rules_path: str | Path,
        registry_path: str | Path | None = None,
        literature_path: str | Path | None = None,
        seed: int | None = None,
        min_k: int = 5,
        max_k: int = 25,
        seed_prob: float = 0.30,
        seed_perturbation: float = 0.05,
    ) -> None:
        self.genre = genre
        self.rules_path = Path(rules_path)
        self.seed = seed
        self.rng = random.Random(seed)
        self.min_k = min_k
        self.max_k = max_k
        self.seed_prob = seed_prob
        self.seed_perturbation = seed_perturbation

        self.rules = json.loads(self.rules_path.read_text())
        if genre not in self.rules["genres"] and genre != "uniform":
            raise ValueError(f"Unknown genre: {genre}")
        self.genre_rules = self.rules["genres"].get(genre, {})
        self.solvent_cas = self.rules["solvent_cas"]

        self.registry = AromaRegistry(registry_path)
        self.compounds = self._load_compounds()
        self.tiers = self._classify_tiers()
        self.active_pool = self._build_active_pool()
        self.seeds = self._load_literature_seeds(literature_path)

    def _load_compounds(self) -> dict[str, dict[str, Any]]:
        """Load all active aroma compounds from the registry."""
        rows = self.registry._conn.execute(
            "SELECT cas, name, smiles, molecular_weight, vapor_pressure_pa, logp FROM aroma_chemicals"
        ).fetchall()
        compounds = {}
        for row in rows:
            record = dict(row)
            cas = record["cas"]
            if cas in self.rules.get("prohibited_cas", []):
                continue
            if cas == self.solvent_cas:
                continue
            # Skip rows lacking a molecular weight: they cannot be tier-classified
            # (MW cutoffs) and would crash comparisons downstream.
            if record.get("molecular_weight") is None:
                continue
            # Assign default guardrails based on IFRA.
            max_w = self._default_max_weight(record)
            record["max_weight_fraction"] = max_w
            record["min_weight_fraction"] = self.rules["default_min_weight"]
            compounds[cas] = record
        return compounds

    def _default_max_weight(self, record: dict[str, Any]) -> float:
        """Assign a default maximum weight fraction based on IFRA."""
        cas = record["cas"]
        if cas in IFRA_RESTRICTIONS:
            limit = IFRA_RESTRICTIONS[cas]["category_4_pct"]
            if limit == 0.0:
                return 0.0
            return min(limit / 100.0, self.rules["default_max_weight"])
        return self.rules["default_max_weight"]

    def _classify_tiers(self) -> dict[str, list[str]]:
        """Classify compounds into top/heart/base tiers by molecular weight."""
        cutoffs = self.rules["tier_cutoffs"]
        tiers: dict[str, list[str]] = {"top": [], "heart": [], "base": []}
        for cas, meta in self.compounds.items():
            mw = meta["molecular_weight"]
            if mw <= cutoffs["top"]["max_mw"]:
                tiers["top"].append(cas)
            elif mw < cutoffs["base"].get("min_mw", 1e9):
                tiers["heart"].append(cas)
            else:
                tiers["base"].append(cas)
        return tiers

    def _build_active_pool(self) -> list[str]:
        """Build the allowed sampling pool for the current genre."""
        if self.genre == "uniform" or self.genre_rules.get("uniform"):
            return [cas for cas, meta in self.compounds.items() if meta["max_weight_fraction"] > 0.0]
        # Genre mode: start with required tiers plus any tier referenced by a
        # min_X_mass constraint, then allow heart/base as supporting elements.
        needed_tiers: set[str] = set(self.genre_rules.get("required_tiers", []))
        for tier in ["top", "heart", "base"]:
            if f"min_{tier}_mass" in self.genre_rules:
                needed_tiers.add(tier)
        needed_tiers.update(["heart", "base"])
        pool: set[str] = set()
        for tier in needed_tiers:
            pool.update(self.tiers.get(tier, []))
        return [cas for cas in pool if self.compounds[cas]["max_weight_fraction"] > 0.0]

    def _load_literature_seeds(self, literature_path: str | Path | None) -> list[dict[str, Any]]:
        """Load literature formulas that are marked as seeds for this genre."""
        if literature_path is None:
            return []
        path = Path(literature_path)
        if not path.exists():
            logger.warning("Literature blueprint not found: %s", path)
            return []
        try:
            formulas = json.loads(path.read_text())
        except Exception as exc:
            logger.warning("Failed to parse literature blueprint: %s", exc)
            return []
        seeds = []
        for f in formulas:
            if f.get("pipeline_role") != "seed":
                continue
            if f.get("expected_profile") != self.genre:
                continue
            components = [c for c in f.get("components", []) if c.get("cas") and c.get("weight_fraction") is not None]
            if not components:
                continue
            # Validate every CAS is present in the registry and pool.
            if any(c["cas"] not in self.compounds for c in components):
                continue
            seeds.append({"formula_id": f.get("formula_id"), "name": f.get("name"), "components": components})
        if seeds:
            logger.info("Loaded %d literature seeds for genre %s", len(seeds), self.genre)
        return seeds

    def _perturb_seed(self, seed_components: list[dict[str, Any]]) -> dict[str, float]:
        """Apply Gaussian noise to seed weights, clip, and renormalize to a unit simplex."""
        weights: dict[str, float] = {}
        for c in seed_components:
            noise = self.rng.gauss(0, self.seed_perturbation)
            weights[c["cas"]] = max(0.0, c["weight_fraction"] + noise)
        total = sum(weights.values())
        if total <= 0.0:
            return weights
        return {cas: w / total for cas, w in weights.items()}

    def _genre_formula_id(self, idx: int) -> str:
        prefix = self.genre[:4].lower()
        return f"gen_{prefix}_{idx:04d}"

    def generate(self, idx: int = 0) -> tuple[list[dict[str, Any]], str]:
        """Generate one candidate formula and return it with a deterministic genre ID."""
        formula = self._generate_formula()
        return formula, self._genre_formula_id(idx)

    def _generate_formula(self) -> list[dict[str, Any]]:
        """Generate a raw {cas, weight_fraction} formula list satisfying genre constraints."""
        solvent_min = self.rules.get("solvent_min", 0.50)
        solvent_max = self.rules.get("solvent_max", 0.95)

        for _ in range(100):
            # Literature-seed path: with seed_prob, perturb a seed formula for this genre.
            if self.seeds and self.rng.random() < self.seed_prob:
                seed = self.rng.choice(self.seeds)
                weights = self._perturb_seed(seed["components"])
                if weights:
                    # Ensure weights respect compound-specific maxima and IFRA limits.
                    weights = self._clip_weights_to_limits(weights)
                    target_aroma = self.rng.uniform(1.0 - solvent_max, 1.0 - solvent_min)
                    scaled = self._project_weights(weights, target_aroma)
                    solvent_share = 1.0 - sum(scaled.values())
                    formula = [{"cas": self.solvent_cas, "weight_fraction": solvent_share}]
                    for cas, w in scaled.items():
                        if w > 0:
                            formula.append({"cas": cas, "weight_fraction": w})
                    return formula

            if self.genre == "uniform" or self.genre_rules.get("uniform"):
                k = self.rng.randint(self.min_k, min(self.max_k, len(self.active_pool)))
                selected = self.rng.sample(self.active_pool, k)
            else:
                selected = self._select_with_tier_constraints()

            if not selected:
                selected = self.rng.sample(self.active_pool, min(self.min_k, len(self.active_pool)))

            target_aroma = self.rng.uniform(1.0 - solvent_max, 1.0 - solvent_min)
            weights = self._assign_tier_constrained_weights(selected, target_aroma)
            if weights:
                solvent_share = 1.0 - sum(weights.values())
                formula = [{"cas": self.solvent_cas, "weight_fraction": solvent_share}]
                for cas, w in weights.items():
                    if w > 0:
                        formula.append({"cas": cas, "weight_fraction": w})
                return formula

        logger.warning("Genre %s constraints infeasible after 100 attempts; falling back", self.genre)
        selected = self.rng.sample(self.active_pool, min(self.min_k, len(self.active_pool)))
        target_aroma = self.rng.uniform(1.0 - solvent_max, 1.0 - solvent_min)
        weights = {cas: self._random_weight(cas) for cas in selected}
        weights = self._project_weights(weights, target_aroma)
        solvent_share = 1.0 - sum(weights.values())
        formula = [{"cas": self.solvent_cas, "weight_fraction": solvent_share}]
        for cas, w in weights.items():
            if w > 0:
                formula.append({"cas": cas, "weight_fraction": w})
        return formula

    def _clip_weights_to_limits(self, weights: dict[str, float]) -> dict[str, float]:
        """Clip each weight to [min_weight_fraction, max_weight_fraction] and renormalize."""
        clipped: dict[str, float] = {}
        for cas, w in weights.items():
            meta = self.compounds.get(cas, {})
            min_w = float(meta.get("min_weight_fraction", 0.001))
            max_w = float(meta.get("max_weight_fraction", 1.0))
            clipped[cas] = max(min_w, min(w, max_w))
        total = sum(clipped.values())
        if total <= 0.0:
            return clipped
        return {cas: w / total for cas, w in clipped.items()}

    def _select_with_tier_constraints(self) -> list[str]:
        """Pick aroma compounds that can physically satisfy the genre tier-mass rules."""
        default_max = self.rules["default_max_weight"]
        selected: set[str] = set()

        # Determine which tiers are needed: explicit required tiers plus any tier
        # that has a min_X_mass constraint.
        needed_tiers: set[str] = set(self.genre_rules.get("required_tiers", []))
        for tier in ["top", "heart", "base"]:
            if f"min_{tier}_mass" in self.genre_rules:
                needed_tiers.add(tier)

        # Ensure at least one compound from each needed tier.
        for tier in needed_tiers:
            candidates = [cas for cas in self.tiers.get(tier, []) if cas in self.active_pool]
            if candidates:
                selected.add(self.rng.choice(candidates))

        # Ensure enough compounds from each constrained tier to satisfy proportions
        # at the largest possible aroma mass (so per-member minimum stays below max).
        max_aroma = 1.0 - self.rules.get("solvent_min", 0.50)
        for tier in ["top", "heart", "base"]:
            key = f"min_{tier}_mass"
            if key not in self.genre_rules:
                continue
            need_mass = self.genre_rules[key] * max_aroma
            present = sum(1 for c in selected if c in self.tiers.get(tier, []))
            need = max(0, int(need_mass / default_max) - present + 1)
            candidates = [cas for cas in self.tiers.get(tier, []) if cas in self.active_pool and cas not in selected]
            need = min(need, len(candidates))
            if need > 0:
                selected.update(self.rng.sample(candidates, need))

        # Fill remaining slots for diversity.
        remaining_slots = self.rng.randint(
            max(0, self.min_k - len(selected)),
            max(0, self.max_k - len(selected)),
        )
        candidates = [cas for cas in self.active_pool if cas not in selected]
        remaining_slots = min(remaining_slots, len(candidates))
        if remaining_slots > 0:
            selected.update(self.rng.sample(candidates, remaining_slots))

        return list(selected)

    def _assign_tier_constrained_weights(
        self,
        selected: list[str],
        target_aroma: float,
    ) -> dict[str, float] | None:
        """
        Assign weights such that tier proportions relative to the aroma mass
        are satisfied. Returns None if infeasible.
        """
        mins = {cas: self.compounds[cas]["min_weight_fraction"] for cas in selected}
        maxs = {cas: self.compounds[cas]["max_weight_fraction"] for cas in selected}

        if self.genre != "uniform" and not self.genre_rules.get("uniform"):
            for tier in ["top", "heart", "base"]:
                key = f"min_{tier}_mass"
                if key not in self.genre_rules:
                    continue
                need = self.genre_rules[key] * target_aroma
                tier_members = [cas for cas in selected if cas in self.tiers.get(tier, [])]
                if not tier_members:
                    return None
                if sum(maxs[cas] for cas in tier_members) < need:
                    return None
                per_member = need / len(tier_members)
                for cas in tier_members:
                    mins[cas] = min(maxs[cas], max(mins[cas], per_member))

            max_top = self.genre_rules.get("max_top_mass")
            if max_top is not None:
                top_members = [cas for cas in selected if cas in self.tiers.get("top", [])]
                if top_members:
                    top_total_max = max_top * target_aroma
                    current_top_max = sum(maxs[cas] for cas in top_members)
                    if current_top_max > top_total_max:
                        scale = top_total_max / current_top_max
                        for cas in top_members:
                            maxs[cas] = max(mins[cas], maxs[cas] * scale)

        if sum(mins.values()) > target_aroma or sum(maxs.values()) < target_aroma:
            return None

        weights = {cas: self.rng.uniform(mins[cas], maxs[cas]) for cas in selected}
        return self._project_weights(weights, target_aroma, mins, maxs)

    def _random_weight(self, cas: str) -> float:
        """Sample a random weight fraction within empirical guardrails."""
        meta = self.compounds[cas]
        min_w = float(meta.get("min_weight_fraction", 0.001))
        max_w = float(meta["max_weight_fraction"])
        if max_w <= min_w:
            return 0.0
        return self.rng.uniform(min_w, max_w)

    def _project_weights(
        self,
        weights: dict[str, float],
        target: float,
        mins: dict[str, float] | None = None,
        maxs: dict[str, float] | None = None,
        tol: float = 1e-12,
    ) -> dict[str, float]:
        """Project weights onto a box-constrained simplex via Lagrange multiplier."""
        cas_list = list(weights.keys())
        x = np.array([weights[cas] for cas in cas_list], dtype=float)
        l = np.array(
            [float((mins or {}).get(cas, self.compounds[cas]["min_weight_fraction"])) for cas in cas_list],
            dtype=float,
        )
        u = np.array(
            [float((maxs or {}).get(cas, self.compounds[cas]["max_weight_fraction"])) for cas in cas_list],
            dtype=float,
        )

        if l.sum() > target or u.sum() < target:
            y = np.clip(x, l, u)
            return dict(zip(cas_list, y))

        def _sum_y(lambda_: float) -> float:
            return float(np.clip(x - lambda_, l, u).sum())

        lambda_low, lambda_high = -1.0, 1.0
        lambda_mid = 0.0
        while _sum_y(lambda_low) < target:
            lambda_low *= 2.0
            if lambda_low < -1e12:
                break
        while _sum_y(lambda_high) > target:
            lambda_high *= 2.0
            if lambda_high > 1e12:
                break

        for _ in range(64):
            lambda_mid = (lambda_low + lambda_high) / 2.0
            s = _sum_y(lambda_mid)
            if abs(s - target) < tol:
                break
            if s > target:
                lambda_low = lambda_mid
            else:
                lambda_high = lambda_mid

        y = np.clip(x - lambda_mid, l, u)
        return dict(zip(cas_list, y))

    def light_ifra_check(self, 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}