src/pino/optimizer.py

from __future__ import annotations

"""
Generative inverse-design layer for PINO.

This module couples the agentic CLI/Space critic with a heuristic search algorithm
(CMA-ES) to evolve manufacture-ready fragrance formulas. The optimizer acts as the
"Composer"; the trained PIMT transformer and physical verifier act as the "Rigid
Critic". Every candidate is normalized to unit mass, clipped to IFRA Category 4
limits, and rejected if it violates physical applicability or safety boundaries.

The intended workflow:

    1. Accept a target sensory profile (objective trajectory or subjective
       psychometrics, or both).
    2. Initialise a population of weight-fraction vectors over a palette drawn from
       the SQLite aroma registry.
    3. Run CMA-ES mutations; after each proposal, enforce:
       - non-negative weights that sum to 1.0,
       - IFRA Category 4 concentration limits,
       - optional renewable/synthetic ingredient restrictions.
    4. Run the local physical verifier to obtain a dry-down trajectory.
    5. Encode the candidate into token/physics tensors and query the PIMT critic.
    6. Score the candidate by inverse MSE against the target profile.
    7. Return the best compliant formula and its predicted trajectory.

All public functions and classes carry docstrings; design choices are explained in
comment blocks where the math or the boundary logic is non-trivial.
"""

import json
import logging
import random
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

from .ingest_formulas import (
    load_literature_manifest,
    normalize_and_unpack_recipe,
    recipe_to_formula_records,
)

import numpy as np
import torch
from cma import CMAEvolutionStrategy

from .embeddings import OlfactoryEmbeddingEngine
from .ifra import IFRA_RESTRICTIONS, check_ifra_restrictions
from .inference_cloud import predict_cloud
from .models import Formula, Ingredient, Status
from .pimt_model import FragranceTrajectoryDataset
from .semantics import get_objective_targets, get_psychometric_targets
from .structures import IngredientResolver
from .verifier import FragrancePipelineVerifier

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

# Curated seed palettes by strategy. These are materials that must be present in
# the search space so the optimizer can satisfy genre-specific briefs (e.g., a
# citrus cologne needs citrus top-note materials). The seed list is resolved
# from the standard materials library or registry; the rest of the palette is
# filled with random registry candidates.
STRATEGY_SEED_CAS: dict[str, list[str]] = {
    "citrus_cologne": [
        "5989-27-5",   # d-Limonene (primary citrus top note)
        "5392-40-5",   # Citral (lemon/citrus top)
        "99-85-4",     # Gamma-terpinene (citrus peel)
        "80-56-8",     # alpha-Pinene (fresh pine/citrus lift)
        "115-95-7",    # Linalyl acetate (bergamet/citrus floral)
        "78-70-6",     # Linalool (light citrus/floral bridge)
    ],
    "fougere": [
        "78-70-6",     # Linalool
        "91-64-5",     # Coumarin
        "121-33-5",    # Vanillin
        "80-56-8",     # alpha-Pinene
        "5989-27-5",   # d-Limonene
    ],
    "floral_woody": [
        "78-70-6",     # Linalool
        "106-24-5",    # Geraniol
        "106-25-2",    # Nerol
        "77-53-2",     # Cedrol
        "5989-27-5",   # d-Limonene
        "121-33-5",    # Vanillin
    ],
    "amber_oriental": [
        "105-95-3",    # Ethylene brassylate
        "91-64-5",     # Coumarin
        "121-33-5",    # Vanillin
        "77-53-2",     # Cedrol
        "8000-27-9",   # Cedarwood oil
    ],
    "wildcard": [],
}

# Minimum weight fraction for an active ingredient to be considered present.
# Anything below this is treated as zero, which keeps the formula sparse and
# manufacture-ready.
MIN_WEIGHT_FRACTION = 1e-4

# Default target for the dry-down trajectory: 8 hours sampled every 10 minutes.
# This matches the synthetic dataset used to train the PIMT model.
DEFAULT_DURATION_S = 8 * 3600.0
DEFAULT_INTERVAL_S = 600.0


@dataclass
class OptimizerConfig:
    """Hyper-parameters for the evolutionary composer."""

    # Search budget. CMA-ES typically needs ~100-300 evaluations per dimension
    # to converge on a continuous constrained problem; 500 is a safe default for
    # palettes up to ~20 ingredients.
    max_iterations: int = 500

    # Population size multiplier. CMAEvolutionStrategy defaults to 4 + 3*log(N);
    # we keep the library default by passing None, but expose it for tuning.
    popsize: int | None = None

    # Initial standard deviation for the CMA-ES search distribution. A value of
    # 0.15 means the first proposals vary weight fractions by roughly ±15% around
    # the initial guess.
    sigma0: float = 0.15

    # Dry-down simulation horizon and sampling interval.
    duration_seconds: float = DEFAULT_DURATION_S
    interval_seconds: float = DEFAULT_INTERVAL_S

    # Penalty for candidates that fail IFRA or physical verification. A large
    # finite value keeps CMA-ES from treating infeasible candidates as attractive
    # while still differentiating them (worse than any feasible MSE).
    infeasible_fitness: float = 1e6

    # Minimum weight fraction for an ingredient to be considered present.
    min_weight_fraction: float = MIN_WEIGHT_FRACTION

    # If True, restrict the palette to materials that can be sourced from
    # renewable feedstocks. This is a coarse heuristic: we exclude synthetic
    # musks, phthalates, and materials explicitly flagged as non-renewable in the
    # IFRA table. The flag is optional and conservative.
    allow_synthetic: bool = True

    # Random seed for reproducibility.
    seed: int | None = None

    # Path to the SQLite aroma registry. If None, the default sibling
    # ``src/pino/registry.db`` is used.
    registry_path: str | Path | None = None

    # Model repository or Space endpoint used by the inference client. None means
    # use the default Space endpoint (preferred for cloud inference).
    model_repo: str | None = None

    # HF token for private Space/Hub access. None means read from HF_TOKEN env.
    token: str | None = None

    # Verbosity: print real-time fitness progression to stderr when > 0.
    verbose: int = 1


@dataclass
class CompositionTarget:
    """
    Target sensory profile used to drive the optimizer.

    The target can encode either or both of:
      - objective: a (T, 138) dry-down trajectory of descriptor intensities.
      - subjective: a 7-dimensional psychometric vector
        [seasonality_4, gender_1, wearability_2].

    Weights control the relative importance of each term in the fitness
    function. By default the objective trajectory dominates, with subjective
    psychometrics providing a secondary signal.
    """

    objective: np.ndarray | None = None
    subjective: np.ndarray | None = None
    weight_objective: float = 1.0
    weight_subjective: float = 0.5
    name: str = "custom"

    strategy: str = "wildcard"

    @staticmethod
    def _infer_strategy(text: str) -> str:
        text_lower = text.lower()
        if "citrus" in text_lower or "summer" in text_lower or "fresh" in text_lower:
            return "citrus_cologne"
        if "amber" in text_lower or "oriental" in text_lower or "winter" in text_lower:
            return "amber_oriental"
        if "floral" in text_lower or "woody" in text_lower:
            return "floral_woody"
        if "fougere" in text_lower or "lavender" in text_lower:
            return "fougere"
        return "wildcard"

    @staticmethod
    def from_text(text: str, *, weight_objective: float = 1.0, weight_subjective: float = 0.5) -> "CompositionTarget":
        """
        Parse a free-form brief into a structured target.

        Examples of supported text:
            - "citrus summer masculine" -> boosts citrus descriptors, summer
              seasonality, and a positive gender profile.
            - "winter amber unisex" -> winter seasonality, amber descriptors,
              neutral gender.
        """
        text_lower = text.lower()
        strategy = CompositionTarget._infer_strategy(text)

        subjective = get_psychometric_targets(strategy)

        # Build an objective trajectory target that reinforces the brief's genre.
        # 49 timesteps = 8 hours sampled every 10 minutes, matching the verifier.
        objective = get_objective_targets(strategy, timesteps=49)

        # Apply crude gender override from the brief.
        if "masculine" in text_lower:
            subjective[4] = 0.7
        elif "feminine" in text_lower:
            subjective[4] = -0.7
        elif "unisex" in text_lower:
            subjective[4] = 0.0

        # Wearability hints: "day" / "office" vs "evening" / "formal".
        if "day" in text_lower or "office" in text_lower or "casual" in text_lower:
            subjective[5] = 0.8
            subjective[6] = 0.2
        elif "evening" in text_lower or "formal" in text_lower or "night" in text_lower:
            subjective[5] = 0.2
            subjective[6] = 0.8

        return CompositionTarget(
            objective=objective,
            subjective=subjective,
            weight_objective=weight_objective,
            weight_subjective=weight_subjective,
            name=text,
            strategy=strategy,
        )

    @staticmethod
    def from_json(path: str | Path) -> "CompositionTarget":
        """Load a target from a JSON brief."""
        with Path(path).open("r", encoding="utf-8") as f:
            payload = json.load(f)

        objective = None
        if "objective" in payload:
            objective = np.array(payload["objective"], dtype=np.float32)

        subjective = None
        if "subjective" in payload:
            subjective = np.array(payload["subjective"], dtype=np.float32)
        elif "brief" in payload:
            return CompositionTarget.from_text(payload["brief"])

        return CompositionTarget(
            objective=objective,
            subjective=subjective,
            weight_objective=payload.get("weight_objective", 1.0),
            weight_subjective=payload.get("weight_subjective", 0.5),
            name=payload.get("name", "custom"),
        )


@dataclass
class Candidate:
    """One evaluated candidate formula plus its metadata."""

    ingredients: list[Ingredient]
    weight_fractions: np.ndarray
    status: str  # "passed", "rejected", "depleted"
    ifra_report: dict[str, Any]
    fitness: float
    prediction: dict[str, Any] = field(default_factory=dict)
    trajectory: list[dict[str, Any]] = field(default_factory=list)
    message: str = ""

    def to_formula_dict(self) -> list[dict[str, Any]]:
        """Return a manufacture-ready list of {cas, weight_fraction} records."""
        return [
            {
                "cas": ing.cas,
                "name": ing.name,
                "weight_fraction": float(wf),
            }
            for ing, wf in zip(self.ingredients, self.weight_fractions)
            if wf >= MIN_WEIGHT_FRACTION
        ]


class FormulationOptimizer:
    """
    Directed evolutionary search for fragrance formulas.

    Uses CMA-ES to optimise continuous weight fractions over a fixed palette of
    ingredients drawn from the ``AromaRegistry``. Every candidate is passed
    through the PINO physical verifier and IFRA Category 4 restrictions before
    it is scored by the PIMT critic.

    Why CMA-ES instead of a generative model?
    -----------------------------------------
    Generative models (VAEs, GANs, diffusion) are excellent at sampling smooth
    manifolds, but fragrance formulation is a constrained combinatorial problem:
    weight fractions must sum to 1.0, IFRA limits are hard thresholds, and the
    physical simulator is expensive and can fail for incompatible combinations.
    CMA-ES handles these constraints via rejection/projection, adapts its search
    covariance to the local landscape, and naturally supports sparse palettes.
    """

    def __init__(self, config: OptimizerConfig | None = None) -> None:
        self.config = config or OptimizerConfig()
        self.resolver = IngredientResolver()
        self.embedding_engine = OlfactoryEmbeddingEngine(use_fallback=True)
        self.verifier = FragrancePipelineVerifier(
            temperature_k=298.15,
            ambient_pressure_pa=101325.0,
            headspace_volume_m3=1e-3,
            liquid_volume_m3=1e-6,
            density_g_ml=0.9,
            surface_area_m2=1e-4,
            mass_transfer_coefficient=1e-4,
        )
        # Cache the PIMT model locally so the composer does not download it for
        # every candidate evaluation. This dramatically speeds up the local fallback.
        self._cached_model = self._load_cached_model()
        self.palette: list[Ingredient] = []
        self.palette_cas: list[str] = []
        self._rng = np.random.default_rng(self.config.seed)
        if self.config.seed is not None:
            random.seed(self.config.seed)
            torch.manual_seed(self.config.seed)

    def _load_cached_model(self) -> "PhysicsInformedMixtureTransformer":
        """Load the critic model once; used for every candidate evaluation."""
        from .pimt_model_hf import PhysicsInformedMixtureTransformer
        logger.info("Loading PIMT critic model from mattbitzesty/pino-pimt")
        model = PhysicsInformedMixtureTransformer.from_pretrained("mattbitzesty/pino-pimt")
        model.eval()
        return model

    def _initial_guess(self, strategy: str, n: int) -> np.ndarray:
        """Return a strategy-biased starting point on the simplex."""
        x = np.zeros(n, dtype=float)
        seed_cas = set(STRATEGY_SEED_CAS.get(strategy, []))
        seed_indices = [i for i, ing in enumerate(self.palette) if ing.cas in seed_cas]
        if not seed_indices:
            return np.ones(n) / n
        # Distribute 70% of the mass evenly across the seeds, 30% across the rest.
        x[seed_indices] = 0.70 / len(seed_indices)
        others = [i for i in range(n) if i not in seed_indices]
        if others:
            x[others] = 0.30 / len(others)
        return x

    def _strategy_bonus(self, weights: np.ndarray, strategy: str) -> float:
        """
        Penalise candidates that lack the brief's expected seed materials.
        This is a soft nudge: missing a large fraction of the expected materials
        adds up to 0.05 to the MSE, enough to break ties but not dominate the
        perceptual loss.
        """
        seed_cas = set(STRATEGY_SEED_CAS.get(strategy, []))
        if not seed_cas:
            return 0.0
        seed_indices = [i for i, ing in enumerate(self.palette) if ing.cas in seed_cas]
        if not seed_indices:
            return 0.0
        seed_mass = float(weights[seed_indices].sum())
        # Missing seed mass penalised quadratically; full seed mass = 0 bonus.
        return 0.20 * (1.0 - seed_mass) ** 2

    def _build_palette(self, count: int = 20, strategy: str = "wildcard") -> None:
        """
        Select a diverse palette from the SQLite registry, seeded with strategy-appropriate
        materials so the optimizer can satisfy genre-specific briefs.
        """
        import sqlite3

        path = Path(self.config.registry_path) if self.config.registry_path else Path(__file__).with_name("registry.db")
        conn = sqlite3.connect(path)
        conn.row_factory = sqlite3.Row

        # Pull candidates with reasonable volatility and a canonical SMILES.
        non_formula_solvents = {"64-17-5", "67-63-0", "25265-71-8"}  # ethanol, IPA, DPG
        restricted_or_problematic = {
            "80-54-6", "83-66-9", "81-15-2", "89-68-9", "92-48-8", "143-50-0"
        }
        strategy_seed = set(STRATEGY_SEED_CAS.get(strategy, []))
        rows = conn.execute(
            """
            SELECT cas, name, smiles, molecular_weight, boiling_point_k, vapor_pressure_pa, logp
            FROM aroma_chemicals
            WHERE molecular_weight BETWEEN 80 AND 300
              AND vapor_pressure_pa IS NOT NULL
              AND vapor_pressure_pa > 0
            ORDER BY molecular_weight ASC
            LIMIT ?
            """,
            (count * 8,),
        ).fetchall()
        conn.close()

        # Start with strategy-seeded materials so the brief has a reachable target.
        seeded: list[Ingredient] = []
        for cas in strategy_seed:
            try:
                ing = self.resolver.resolve(cas, weight_fraction=0.0)
                seeded.append(ing)
            except Exception as exc:
                logger.debug("Could not resolve seed material %s: %s", cas, exc)

        candidates: list[Ingredient] = []
        for row in rows:
            cas = str(row["cas"])
            if cas in non_formula_solvents or cas in restricted_or_problematic or cas in strategy_seed:
                continue
            try:
                ing = self.resolver.resolve(
                    cas,
                    weight_fraction=0.0,
                    name=row["name"],
                )
                candidates.append(ing)
            except Exception as exc:
                logger.debug("Skipping palette candidate %s: %s", cas, exc)
                continue

        # Fill the rest of the palette with random registry candidates, preserving
        # the seeds. We always want at least the seeds present.
        remaining = max(0, count - len(seeded))
        if remaining > 0 and candidates:
            indices = self._rng.choice(len(candidates), size=min(remaining, len(candidates)), replace=False)
            self.palette = seeded + [candidates[i] for i in indices]
        else:
            self.palette = seeded
        self.palette_cas = [ing.cas for ing in self.palette]
        logger.info(
            "Built palette of %d ingredients (%d seeded) from %d registry candidates",
            len(self.palette),
            len(seeded),
            len(candidates),
        )

    def _restrict_palette(self) -> None:
        """Apply the renewable/synthetic restriction if configured."""
        if self.config.allow_synthetic:
            return

        non_renewable = {
            "84-66-2",   # Diethyl phthalate
            "131-11-3",  # Dimethyl phthalate
            "81-14-1",   # Musk ketone
            "83-66-9",   # Musk ambrette
            "81-15-2",   # Musk xylene
            "105-95-3",  # Ethylene brassylate
        }
        self.palette = [ing for ing in self.palette if ing.cas not in non_renewable]
        self.palette_cas = [ing.cas for ing in self.palette]
        logger.info("Restricted palette to %d renewable ingredients", len(self.palette))

    def _project_and_normalize(self, x: np.ndarray) -> np.ndarray:
        """
        Project a raw CMA-ES vector to a feasible weight-fraction simplex.

        Steps:
          1. Clip to [0, 1] and zero out negligible components (promotes sparsity).
          2. Enforce IFRA Category 4 limits by capping each ingredient at its
             maximum allowed concentration. Because re-normalisation can push a
             capped ingredient back above its limit, we iterate the cap-and-rescale
             step until the simplex is stable (water-filling on the simplex).
          3. Re-normalize so the vector sums to exactly 1.0.

        This projection is cheap and keeps the search inside the safe region.
        """
        x = np.clip(x, 0.0, 1.0)
        x[x < self.config.min_weight_fraction] = 0.0

        # Build the limit vector. If an ingredient has no IFRA entry, its limit
        # is 1.0 (i.e., it can be the entire formula). Prohibited materials have
        # limit 0.0 and are therefore removed by this step.
        limits = np.ones(len(self.palette), dtype=float)
        for i, ing in enumerate(self.palette):
            restriction = IFRA_RESTRICTIONS.get(ing.cas)
            if restriction:
                limits[i] = restriction["category_4_pct"] / 100.0

        # Iterative water-filling on the simplex. Any component above its limit
        # is clamped to the limit; the excess mass is redistributed among the
        # components that are still below their limits. This terminates quickly
        # because at least one component is capped each iteration.
        for _ in range(50):
            x = np.minimum(x, limits)
            total = x.sum()
            if total <= 0:
                break
            scale = 1.0 / total
            x = x * scale
            # If all components are now within limits after rescaling, we are done.
            if np.all(x <= limits + 1e-12):
                break

        total = x.sum()
        if total > 0:
            x = x / total
        else:
            # Fallback: uniform distribution over ingredients with a non-zero limit.
            allowed = limits > 0
            x = np.zeros_like(x)
            x[allowed] = 1.0 / allowed.sum()
        return x

    def _evaluate(self, x: np.ndarray, target: CompositionTarget) -> Candidate:
        """
        Evaluate one candidate weight vector against the target.

        The pipeline is intentionally defensive:
          1. Project to a feasible simplex.
          2. Build a minimal formula record and run the physical verifier.
          3. If the formula is rejected or depleted, return a penalised fitness.
          4. Otherwise encode the candidate for the PIMT critic and compute MSE.
          5. Return the candidate with its prediction and fitness.
        """
        weights = self._project_and_normalize(x)
        formula_records = [
            {"cas": ing.cas, "weight_fraction": float(wf)}
            for ing, wf in zip(self.palette, weights)
            if wf >= self.config.min_weight_fraction
        ]

        if len(formula_records) < 2:
            return Candidate(
                ingredients=self.palette,
                weight_fractions=weights,
                status="rejected",
                ifra_report={"passed": False, "violations": [], "max_usage": []},
                fitness=self.config.infeasible_fitness,
                message="Too few ingredients",
            )

        # Physical / IFRA verification.
        result = self.verifier.run_sim(
            formula_records,
            duration_seconds=self.config.duration_seconds,
            interval_seconds=self.config.interval_seconds,
        )
        if result["status"] in ("rejected", Status.REJECTED):
            return Candidate(
                ingredients=self.palette,
                weight_fractions=weights,
                status="rejected",
                ifra_report=result["ifra_report"],
                fitness=self.config.infeasible_fitness,
                message=result["message"],
            )

        # Encode for the PIMT critic. We use the cached local model directly to
        # avoid the per-request HTTP round-trip to the Space endpoint, which is
        # especially slow when the Space is sleeping or rebuilding.
        try:
            prediction = self._predict_local(formula_records, result["trajectory"])
        except Exception as exc:
            logger.warning("Prediction failed for candidate: %s", exc)
            return Candidate(
                ingredients=self.palette,
                weight_fractions=weights,
                status=result["status"],
                ifra_report=result["ifra_report"],
                fitness=self.config.infeasible_fitness,
                trajectory=result["trajectory"],
                message=f"Prediction error: {exc}",
            )

        mse = self._mse(prediction, target)
        fitness = mse if mse > 0 else 1e-12
        strategy_bonus = self._strategy_bonus(weights, target.strategy)
        fitness = fitness + strategy_bonus

        return Candidate(
            ingredients=self.palette,
            weight_fractions=weights,
            status=result["status"],
            ifra_report=result["ifra_report"],
            fitness=fitness,
            prediction=prediction,
            trajectory=result["trajectory"],
            message=result["message"],
        )

    def _predict(
        self,
        formula_records: list[dict[str, Any]],
        trajectory: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """
        Build a lightweight formula payload and query the PIMT critic.

        We intentionally reuse the cloud inference client so that the composer can
        run against a local CPU fallback or the private HF Space endpoint without
        code changes. For pure local optimisation, this falls back to loading the
        model from ``mattbitzesty/pino-pimt``.
        """
        formula_payload = {
            "formula": formula_records,
            "trajectory": trajectory,
            "metadata": {"generation_strategy": "wildcard"},
        }
        return predict_cloud(formula_payload, token=self.config.token, _cached_model=self._cached_model)

    def _predict_local(
        self,
        formula_records: list[dict[str, Any]],
        trajectory: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Run the cached PIMT critic directly without any HTTP round-trip."""
        import torch
        from .inference_cloud import encode_formula

        formula_payload = {
            "formula": formula_records,
            "trajectory": trajectory,
            "metadata": {"generation_strategy": "wildcard"},
        }
        payload = encode_formula(formula_payload)
        tokens = torch.tensor(payload["tokens"], dtype=torch.float32).unsqueeze(0)
        physics = torch.tensor(payload["physics"], dtype=torch.float32).unsqueeze(0)
        mask = torch.tensor(payload["src_key_padding_mask"], dtype=torch.bool).unsqueeze(0)

        with torch.no_grad():
            output = self._cached_model(tokens, physics, src_key_padding_mask=mask)

        return {
            "status": "success (local cached)",
            "objective": output.logits.cpu().numpy().tolist(),
            "subjective": output.subjective.cpu().numpy().tolist(),
        }

    @staticmethod
    def _compute_olfactory_mse(
        pred_obj: np.ndarray,
        tgt_obj: np.ndarray,
        active_indices: tuple[int, ...] = (0, 1, 2, 3),
        active_weight: float = 50.0,
        background_weight: float = 0.1,
    ) -> float:
        """
        Importance-weighted MSE for the objective trajectory.

        Active genre descriptors (e.g., citrus, lavender, wood, musk) are weighted
        heavily so the optimizer is driven by the brief. Background channels are
        down-weighted so they do not dilute the signal.
        """
        squared_errors = (pred_obj - tgt_obj) ** 2
        weights = np.ones_like(squared_errors) * background_weight
        for idx in active_indices:
            weights[..., idx] = active_weight
        return float(np.mean(squared_errors * weights))

    def _mse(self, prediction: dict[str, Any], target: CompositionTarget) -> float:
        """
        Compute a weighted mean-squared-error loss between prediction and target.
        """
        errors = []

        if target.objective is not None and "objective" in prediction:
            pred_obj = np.array(prediction["objective"], dtype=np.float32)
            if pred_obj.ndim == 3:
                pred_obj = pred_obj[0]
            tgt_obj = target.objective
            if tgt_obj.ndim == 1:
                tgt_obj = np.tile(tgt_obj, (pred_obj.shape[0], 1))
            # Use an importance-weighted MSE: active genre descriptors (0-3) get
            # a heavy weight while the 134 background channels are down-weighted
            # so they cannot drown the brief signal.
            obj_error = self._compute_olfactory_mse(pred_obj, tgt_obj)
            errors.append(target.weight_objective * obj_error)

        if target.subjective is not None and "subjective" in prediction:
            pred_sub = np.array(prediction["subjective"], dtype=np.float32)
            if pred_sub.ndim == 2:
                pred_sub = pred_sub[0]
            sub_error = float(np.mean((pred_sub - target.subjective) ** 2))
            errors.append(target.weight_subjective * sub_error)

        if not errors:
            return self.config.infeasible_fitness
        return float(np.mean(errors))

    def optimize(
        self,
        target: CompositionTarget,
        palette_size: int = 20,
        initial_guess: np.ndarray | None = None,
        progress_callback: Callable[[int, float, Candidate], None] | None = None,
        seed_recipe: dict[str, Any] | None = None,
    ) -> Candidate:
        """
        Run the evolutionary search and return the best compliant candidate.

        Parameters
        ----------
        target: CompositionTarget
            Desired sensory profile.
        palette_size: int
            Number of ingredients to draw from the registry.
        initial_guess: np.ndarray | None
            Optional starting point on the simplex. If None, a uniform mixture is
            used as the CMA-ES centroid.
        progress_callback: callable | None
            Optional function ``(iteration, fitness, best_candidate) -> None`` called
            after each generation. The CLI uses this to print progress to stderr.
        seed_recipe: dict | None
            Optional normalized literature recipe whose CAS/weights seed the first
            generation. The recipe's materials are guaranteed to be in the palette,
            and its weight vector becomes the CMA-ES centroid.
        """
        self._build_palette(palette_size, strategy=target.strategy)
        self._restrict_palette()

        if seed_recipe:
            self._inject_seed_recipe(seed_recipe)

        n = len(self.palette)
        if n == 0:
            raise RuntimeError("Palette is empty; cannot optimise")

        # Strategy-aware initial guess: give a strong prior to the materials that
        # are expected to satisfy the brief. This keeps CMA-ES from getting stuck
        # in a generic floral/heart local minimum.
        if initial_guess is not None:
            x0 = initial_guess
        elif seed_recipe is not None:
            x0 = self.recipe_initial_guess(seed_recipe)
        else:
            x0 = self._initial_guess(target.strategy, n)
        x0 = self._project_and_normalize(x0)
        # CMA-ES minimises the objective; lower MSE = higher fitness.
        es = CMAEvolutionStrategy(
            x0=x0,
            sigma0=self.config.sigma0,
            inopts={"popsize": self.config.popsize, "verbose": -9},
        )

        best: Candidate | None = None
        iteration = 0

        while not es.stop() and iteration < self.config.max_iterations:
            solutions = es.ask()
            fitness_values: list[float] = []

            for x in solutions:
                candidate = self._evaluate(x, target)
                fitness_values.append(candidate.fitness)

                if best is None or candidate.fitness < best.fitness:
                    best = candidate

            es.tell(solutions, fitness_values)
            iteration += 1

            if self.config.verbose > 0 and iteration % 10 == 0:
                logger.info(
                    "Iteration %d / %d: best fitness (MSE) = %.6f",
                    iteration,
                    self.config.max_iterations,
                    best.fitness if best else float("inf"),
                )

            if progress_callback:
                progress_callback(iteration, best.fitness if best else float("inf"), best)

        if best is None:
            raise RuntimeError("Optimization failed to produce any feasible candidate")

        # Re-run the best candidate to capture a full prediction/trajectory if it
        # was produced by an infeasible placeholder during the first iteration.
        best = self._evaluate(best.weight_fractions, target)
        return best

    def _inject_seed_recipe(self, seed_recipe: dict[str, Any]) -> None:
        """
        Ensure every CAS in a normalized literature recipe is present in the palette,
        and patch the palette order so the recipe weight vector aligns with the
        optimizer's internal arrays.
        """
        recipe_cas = {c["cas"] for c in seed_recipe.get("components", [])}
        current_cas = {ing.cas for ing in self.palette}
        missing = recipe_cas - current_cas
        if missing:
            logger.info("Expanding palette to include %d recipe CAS entries", len(missing))
            for cas in missing:
                try:
                    ing = self.resolver.resolve(cas, weight_fraction=0.0)
                    self.palette.append(ing)
                except Exception as exc:
                    logger.warning("Could not resolve seed recipe CAS %s: %s", cas, exc)
        self.palette_cas = [ing.cas for ing in self.palette]

    def recipe_initial_guess(self, seed_recipe: dict[str, Any]) -> np.ndarray:
        """Build a weight vector from a normalized recipe aligned to the current palette."""
        if not self.palette:
            raise RuntimeError("Palette not built before constructing recipe initial guess")
        x = np.zeros(len(self.palette), dtype=float)
        recipe_by_cas = {c["cas"]: c["weight_fraction"] for c in seed_recipe.get("components", [])}
        for i, ing in enumerate(self.palette):
            x[i] = recipe_by_cas.get(ing.cas, 0.0)
        total = x.sum()
        if total <= 0:
            raise RuntimeError("Seed recipe has zero total weight after palette mapping")
        return x / total


# ---------------------------------------------------------------------------
# Convenience helpers for natural-language / JSON brief entry points.
# ---------------------------------------------------------------------------


def compose_from_text(
    brief: str,
    *,
    max_iterations: int = 500,
    allow_synthetic: bool = True,
    palette_size: int = 20,
    verbose: int = 1,
    seed: int | None = None,
) -> dict[str, Any]:
    """
    High-level helper: compose a formula from a natural-language brief.

    Returns a JSON-serialisable dict with the optimised formula, status, fitness,
    and a short human-readable message.
    """
    target = CompositionTarget.from_text(brief)
    config = OptimizerConfig(
        max_iterations=max_iterations,
        allow_synthetic=allow_synthetic,
        seed=seed,
        verbose=verbose,
    )
    optimizer = FormulationOptimizer(config)
    best = optimizer.optimize(target, palette_size=palette_size)

    return {
        "brief": brief,
        "target": target.name,
        "formula": best.to_formula_dict(),
        "status": best.status,
        "fitness": best.fitness,
        "ifra_passed": best.ifra_report.get("passed", False),
        "message": best.message,
    }


def compose_from_json(
    path: str | Path,
    *,
    max_iterations: int = 500,
    allow_synthetic: bool = True,
    palette_size: int = 20,
    verbose: int = 1,
    seed: int | None = None,
) -> dict[str, Any]:
    """High-level helper: compose a formula from a JSON target brief."""
    target = CompositionTarget.from_json(path)
    config = OptimizerConfig(
        max_iterations=max_iterations,
        allow_synthetic=allow_synthetic,
        seed=seed,
        verbose=verbose,
    )
    optimizer = FormulationOptimizer(config)
    best = optimizer.optimize(target, palette_size=palette_size)

    return {
        "target": target.name,
        "formula": best.to_formula_dict(),
        "status": best.status,
        "fitness": best.fitness,
        "ifra_passed": best.ifra_report.get("passed", False),
        "message": best.message,
    }