src/pino/verifier.py

from __future__ import annotations

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

import numpy as np

from .evaporation import EvaporationModel
from .ifra import check_ifra_restrictions
from .models import Formula, Ingredient, IngredientResolutionError, Status, Trajectory
from .structures import IngredientResolver
from .vle import VLECalculator

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


class FragrancePipelineVerifier:
    """
    High-level end-to-end verifier for synthetic fragrance formulations.

    Ingests a list of {cas, weight_fraction} records, resolves chemical
    attributes, checks safety constraints, and runs the stiff ODE evaporation
    model to produce a structured trajectory result.
    """

    def __init__(
        self,
        temperature_k: float = 298.15,
        ambient_pressure_pa: float = 101325.0,
        headspace_volume_m3: float = 1e-3,
        liquid_volume_m3: float = 1e-6,
        density_g_ml: float = 0.9,
        surface_area_m2: float = 1e-4,
        mass_transfer_coefficient: float = 1e-4,
    ) -> None:
        self.temperature_k = temperature_k
        self.ambient_pressure_pa = ambient_pressure_pa
        self.headspace_volume_m3 = headspace_volume_m3
        self.liquid_volume_m3 = liquid_volume_m3
        self.density_g_ml = density_g_ml
        self.surface_area_m2 = surface_area_m2
        self.mass_transfer_coefficient = mass_transfer_coefficient
        self.resolver = IngredientResolver()

    def run_sim(
        self,
        formula: list[dict[str, Any]],
        duration_seconds: float = 3600.0,
        interval_seconds: float = 600.0,
        method: str = "BDF",
        skip_ifra: bool = False,
        skip_unresolved: bool = False,
    ) -> dict[str, Any]:
        """
        Run the full verification pipeline on a draft formula.

        Returns a dict with:
            - status: "passed" or "rejected"
            - trajectory: list of per-time-step records
            - depletion_rates: dict of CAS -> fractional depletion
            - ifra_report: IFRA restriction check
            - applicability_errors: list of out-of-domain errors
            - message: human-readable summary
        """
        # Resolve ingredients. By default, a verifier pass means every supplied
        # component was inside the applicability domain. Dataset builders may
        # opt into partial simulation after preserving unresolved rows upstream.
        try:
            resolved_pairs, errors = self.resolver.resolve_formula_pairs(formula, skip_unresolved=skip_unresolved)
            if errors:
                for exc in errors:
                    logger.warning("Skipping unresolved ingredient: %s", exc)
        except IngredientResolutionError as exc:
            logger.warning("Formula rejected: out of applicability domain for %s", exc.cas)
            return {
                "status": "rejected",
                "trajectory": [],
                "depletion_rates": {},
                "ifra_report": {"passed": False, "violations": [], "max_usage": []},
                "applicability_errors": [
                    {"cas": exc.cas, "smiles": exc.smiles, "reason": str(exc)}
                ],
                "message": f"Out of applicability domain: {exc}",
            }

        if not resolved_pairs:
            return {
                "status": "rejected",
                "trajectory": [],
                "depletion_rates": {},
                "ifra_report": {"passed": False, "violations": [], "max_usage": []},
                "applicability_errors": [],
                "message": "No resolvable ingredients in formula",
            }

        ingredients = [ing for _, ing in resolved_pairs]
        weights = [float(rec.get("weight_fraction", 0.0)) for rec, _ in resolved_pairs]
        # Re-normalize in case skipped ingredients changed total
        total_w = sum(weights)
        if total_w <= 0:
            return {
                "status": "rejected",
                "trajectory": [],
                "depletion_rates": {},
                "ifra_report": {"passed": False, "violations": [], "max_usage": []},
                "applicability_errors": [],
                "message": "Formula has no positive weight fractions",
            }
        weights = [w / total_w for w in weights]

        resolved_formula = Formula(
            ingredients=tuple(ingredients),
            weight_fractions=tuple(weights),
            temperature_k=self.temperature_k,
            ambient_pressure_pa=self.ambient_pressure_pa,
            headspace_volume_m3=self.headspace_volume_m3,
            liquid_volume_m3=self.liquid_volume_m3,
            metadata={"density_g_ml": self.density_g_ml},
        )

        # IFRA check
        ifra_report = check_ifra_restrictions(resolved_formula)
        if not skip_ifra and not ifra_report["passed"]:
            logger.warning("Formula rejected by IFRA restrictions: %s", ifra_report["violations"])
            return {
                "status": "rejected",
                "trajectory": [],
                "depletion_rates": {},
                "ifra_report": ifra_report,
                "applicability_errors": [],
                "message": f"IFRA restriction violation: {ifra_report['violations']}",
            }

        # Run evaporation model
        vle = VLECalculator(resolved_formula)
        model = EvaporationModel(
            resolved_formula,
            vle=vle,
            surface_area_m2=self.surface_area_m2,
            mass_transfer_coefficient=self.mass_transfer_coefficient,
        )
        n_steps = max(2, int(duration_seconds / interval_seconds) + 1)
        t_eval = np.linspace(0.0, duration_seconds, n_steps)
        trajectory = model.solve(
            t_span=(0.0, duration_seconds),
            t_eval=t_eval,
            method=method,
        )

        # Compute depletion rates per CAS
        depletion_rates = self._depletion_rates(trajectory)

        # Convert trajectory to list of dicts
        trajectory_records = self._trajectory_to_records(trajectory)

        # Detect film-depletion terminal event
        if trajectory.t[-1] < duration_seconds:
            status = Status.DEPLETED
            message = f"Film depleted at t={trajectory.t[-1]:.1f}s before horizon {duration_seconds:.1f}s"
        else:
            status = Status.PASSED
            message = "Formula verified successfully"

        return {
            "status": status,
            "trajectory": trajectory_records,
            "depletion_rates": depletion_rates,
            "ifra_report": ifra_report,
            "applicability_errors": [],
            "message": message,
        }

    @staticmethod
    def _depletion_rates(trajectory: Trajectory) -> dict[str, float]:
        """True fractional liquid molar depletion."""
        rates: dict[str, float] = {}
        n0 = trajectory.n_liquid[:, 0]
        n1 = trajectory.n_liquid[:, -1]
        for idx, ing in enumerate(trajectory.formula.ingredients):
            if n0[idx] > 0:
                rates[ing.cas] = float((n0[idx] - n1[idx]) / n0[idx])
            else:
                rates[ing.cas] = 0.0
        return rates

    @staticmethod
    def _trajectory_to_records(trajectory: Trajectory) -> list[dict[str, Any]]:
        """Convert a Trajectory object into a list of JSON-serializable records keyed by CAS."""
        records: list[dict[str, Any]] = []
        species = [ing.cas for ing in trajectory.formula.ingredients]
        for i, t in enumerate(trajectory.t):
            records.append({
                "t_s": float(t),
                "x_liquid": {s: float(trajectory.x_liquid[j, i]) for j, s in enumerate(species)},
                "C_gas_mg_m3": {s: float(trajectory.C_gas[j, i]) for j, s in enumerate(species)},
                "y_gas": {s: float(trajectory.y_gas[j, i]) for j, s in enumerate(species)},
                "OAV": {s: float(trajectory.OAV[j, i]) for j, s in enumerate(species)},
            })
        return records


class FormulaVerifier:
    """Original high-level verifier: intake a draft formula and emit a JSON trajectory."""

    def __init__(self, formula: Formula) -> None:
        self.formula = formula
        self.vle = VLECalculator(formula)
        self.model = EvaporationModel(formula, vle=self.vle)

    def verify(self, **solve_kwargs: Any) -> Trajectory:
        return self.model.solve(**solve_kwargs)

    def verify_and_write(
        self,
        output_path: str | Path = "trajectory_output.json",
        **solve_kwargs: Any,
    ) -> Path:
        trajectory = self.verify(**solve_kwargs)
        path = Path(output_path)
        path.write_text(json.dumps(trajectory.to_dict(), indent=2))
        return path


def verify_formula_dict(data: dict[str, Any], **solve_kwargs: Any) -> Trajectory:
    """Convenience entrypoint: build a Formula from a dict and run the verifier."""
    from .structures import ingredient_from_data

    ingredients = [ingredient_from_data(i) for i in data["ingredients"]]
    formula = Formula(
        ingredients=tuple(ingredients),
        weight_fractions=tuple(data["weight_fractions"]),
        temperature_k=float(data.get("temperature_k", 298.15)),
        ambient_pressure_pa=float(data.get("ambient_pressure_pa", 101325.0)),
        headspace_volume_m3=float(data.get("headspace_volume_m3", 1e-3)),
        liquid_volume_m3=float(data.get("liquid_volume_m3", 1e-6)),
        metadata=data.get("metadata", {}),
    )
    verifier = FormulaVerifier(formula)
    return verifier.verify(**solve_kwargs)