src/pino/models.py

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import numpy as np


class Status(str, Enum):
    PASSED = "passed"
    REJECTED = "rejected"
    DEPLETED = "depleted"


class IngredientResolutionError(ValueError):
    """Raised when a raw ingredient cannot be resolved or fragmented."""

    def __init__(self, message: str, cas: str | None = None, smiles: str | None = None) -> None:
        super().__init__(message)
        self.cas = cas
        self.smiles = smiles


@dataclass(frozen=True, slots=True)
class Ingredient:
    """Canonical representation of one formula component."""

    name: str
    cas: str
    smiles: str
    molecular_weight: float  # g/mol
    vapor_pressure_pa: float  # pure-component saturation pressure at 25 °C
    boiling_point_k: float | None = None
    odor_threshold_ug_m3: float | None = None  # used to compute OAV
    logp: float | None = None  # RDKit Crippen LogP
    # UNIFAC group counts: group_id -> integer count
    unifac_groups: dict[str, int] = field(default_factory=dict)
    # Optional constants for Antoine / DIPPR extension
    antoine_coeffs: tuple[float, float, float] | None = None

    def __post_init__(self) -> None:
        if self.molecular_weight <= 0:
            raise ValueError(f"MW must be positive for {self.name}")
        if self.vapor_pressure_pa <= 0:
            raise ValueError(f"vapor_pressure_pa must be positive for {self.name}")


@dataclass(frozen=True, slots=True)
class Formula:
    """A multi-component fragrance formulation."""

    ingredients: tuple[Ingredient, ...]
    weight_fractions: tuple[float, ...]  # must sum to 1.0
    temperature_k: float = 298.15
    ambient_pressure_pa: float = 101325.0
    headspace_volume_m3: float = 1e-3  # 1 L default
    liquid_volume_m3: float = 1e-6  # 1 mL default
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if len(self.ingredients) != len(self.weight_fractions):
            raise ValueError("ingredients and weight_fractions must have the same length")
        total = sum(self.weight_fractions)
        if not (0.999 <= total <= 1.001):
            raise ValueError(f"weight fractions must sum to 1.0, got {total}")
        if self.temperature_k <= 0 or self.ambient_pressure_pa <= 0:
            raise ValueError("temperature and pressure must be positive")

    @property
    def mole_fractions(self) -> np.ndarray:
        """Convert weight fractions to liquid mole fractions."""
        weights = np.array(self.weight_fractions, dtype=float)
        mws = np.array([i.molecular_weight for i in self.ingredients], dtype=float)
        moles = weights / mws
        return moles / moles.sum()

    @property
    def avg_molecular_weight(self) -> float:
        return float(np.dot(self.mole_fractions, [i.molecular_weight for i in self.ingredients]))


@dataclass(frozen=True, slots=True)
class Trajectory:
    """Time-resolved output of the evaporation simulation."""

    t: np.ndarray
    n_liquid: np.ndarray  # absolute moles remaining (n_species, n_times)
    x_liquid: np.ndarray  # shape (n_species, n_times)
    C_gas: np.ndarray  # shape (n_species, n_times)
    y_gas: np.ndarray  # shape (n_species, n_times)
    OAV: np.ndarray  # shape (n_species, n_times)
    formula: Formula
    notes: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "t_s": self.t.tolist(),
            "species": [i.name for i in self.formula.ingredients],
            "n_liquid_mol": self.n_liquid.tolist(),
            "x_liquid": self.x_liquid.tolist(),
            "C_gas_mg_m3": self.C_gas.tolist(),
            "y_gas_mole_fraction": self.y_gas.tolist(),
            "OAV": self.OAV.tolist(),
            "temperature_k": self.formula.temperature_k,
            "ambient_pressure_pa": self.formula.ambient_pressure_pa,
            "notes": self.notes,
        }