src/pino/evaporation.py

from __future__ import annotations

import logging
from typing import Any

import numpy as np
from scipy.integrate import solve_ivp

from .models import Formula, Trajectory
from .vle import VLECalculator

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


class EvaporationModel:
    """
    Stiff ODE model for multi-component fragrance evaporation.

    State vector: n_i(t) = moles of component i remaining in the liquid film.
    Assumptions:
      - Well-mixed thin liquid film of fixed surface area A (m2).
      - Evaporation flux J_i = k_i * (f_i^L / P) (mol/m2/s) using liquid fugacity ratios.
      - k_i = k_base * (MW_i / MW_ethanol)^alpha / (P_i^sat)^beta as a volatility scaling.
      - Gas phase is instantaneously well-mixed with the headspace volume.
    """

    def __init__(
        self,
        formula: Formula,
        vle: VLECalculator | None = None,
        surface_area_m2: float = 1e-4,  # 1 cm^2 film
        mass_transfer_coefficient: float = 1e-4,  # m/s base
        volatility_exponent: float = 0.5,
        mw_exponent: float = -0.5,
        skin_sink_alpha: float = -6.0,   # Potts-Guy intercept (log10 scale)
        skin_sink_beta: float = 0.4,     # Potts-Guy LogP coefficient
        skin_sink_gamma: float = 0.005,  # Potts-Guy MW coefficient
        skin_sink_max: float = 1e-6,     # s^-1 clamp
    ) -> None:
        self.formula = formula
        self.vle = vle or VLECalculator(formula)
        self.surface_area = surface_area_m2
        self.k_base = mass_transfer_coefficient
        self.alpha = volatility_exponent
        self.beta = mw_exponent
        self.skin_sink_alpha = skin_sink_alpha
        self.skin_sink_beta = skin_sink_beta
        self.skin_sink_gamma = skin_sink_gamma
        self.skin_sink_max = skin_sink_max
        self._mws = np.array([i.molecular_weight for i in formula.ingredients], dtype=float)
        self._p_sat = np.array([i.vapor_pressure_pa for i in formula.ingredients], dtype=float)
        self._logp = np.array(
            [i.logp if i.logp is not None else 0.0 for i in formula.ingredients], dtype=float
        )
        self._k_skin = self._compute_skin_sink_rates()
        self._n0 = self._initial_moles()
        # Terminal film-depletion event for solve_ivp.
        self._film_depleted_event = self._make_film_depleted_event()

    def _make_film_depleted_event(self):
        """Return a solve_ivp event that fires when the total liquid film is depleted."""
        def event(t: float, n: np.ndarray) -> float:
            # Triggers when total liquid moles drops below a physical trace limit.
            return float(np.maximum(n, 0.0).sum() - 1e-7)

        event.terminal = True
        return event

    def _initial_moles(self) -> np.ndarray:
        """Moles of each component in the initial 1 mL liquid sample."""
        # Assume liquid density ~ 0.9 g/mL for fragrance oil (overridden by formula metadata)
        total_volume_l = self.formula.liquid_volume_m3 * 1000.0
        density_g_ml = self.formula.metadata.get("density_g_ml", 0.9)
        total_mass_g = total_volume_l * density_g_ml
        w = np.array(self.formula.weight_fractions, dtype=float)
        masses_g = w * total_mass_g
        return masses_g / self._mws

    def _mass_transfer_coefficients(self, x: np.ndarray) -> np.ndarray:
        """Species-specific k_i."""
        mw_ethanol = 46.07
        scaling = (self._mws / mw_ethanol) ** self.beta
        # Lower saturation pressure => higher boiling point => slower evaporation
        volatility_scaling = (self._p_sat / 1e5) ** self.alpha
        return self.k_base * scaling / (volatility_scaling + 1e-12)

    def _compute_skin_sink_rates(self) -> np.ndarray:
        """
        Potts-Guy inspired per-species skin absorption rate constant (s^-1).
        log10(k_skin) = alpha + beta*logP - gamma*MW
        """
        log10_k = (
            self.skin_sink_alpha
            + self.skin_sink_beta * self._logp
            - self.skin_sink_gamma * self._mws
        )
        with np.errstate(over="ignore", under="ignore"):
            k_skin = 10.0 ** log10_k
        return np.clip(k_skin, 0.0, self.skin_sink_max)

    def _skin_absorption_rate(self, n: np.ndarray) -> np.ndarray:
        """
        First-order skin sink: rate_i = k_skin_i * n_i (mol/s).
        Vanishes as n_i -> 0, so it cannot drive moles negative.
        """
        return self._k_skin * n

    def _surface_coverage(self, n_total: float) -> float:
        """
        Smooth, differentiable surface coverage factor that gates all mass transport
        as the liquid film approaches total exhaustion. Avoids the n_total -> 0
        cliff in the mole-fraction calculation.
        """
        return float(np.tanh(n_total / 1e-6))

    def _dn_dt(self, n: np.ndarray) -> np.ndarray:
        """Physical d n_i / dt = -theta * (evaporation + skin absorption)."""
        # Clamp individual species to a tiny positive floor to keep x_i > 0.
        n = np.maximum(n, 1e-30)
        n_total = n.sum()
        # Smoothly gate the entire derivative as the film dries out.
        theta = self._surface_coverage(n_total)
        if theta <= 1e-12:
            return np.zeros_like(n)
        x = n / n_total
        fugacity_ratio = self.vle.fugacity_ratios(x)
        k_i = self._mass_transfer_coefficients(x)
        # flux (mol/m2/s) * area (m2) = mol/s
        evaporation = -k_i * self.surface_area * fugacity_ratio
        skin_sink = -self._skin_absorption_rate(n)
        return theta * (evaporation + skin_sink)

    def _rhs(self, t: float, n: np.ndarray) -> np.ndarray:
        """d n_i / dt = - evaporation - skin absorption."""
        return self._dn_dt(n)

    def solve(
        self,
        t_span: tuple[float, float] = (0.0, 28800.0),
        t_eval: np.ndarray | None = None,
        method: str = "BDF",
        rtol: float = 1e-6,
        atol: float = 1e-12,
        film_depletion_threshold: float = 1e-6,
    ) -> Trajectory:
        if t_eval is None:
            t_eval = np.linspace(t_span[0], t_span[1], 481)

        res = solve_ivp(
            self._rhs,
            t_span,
            self._n0,
            method=method,
            t_eval=t_eval,
            events=self._film_depleted_event,
            rtol=rtol,
            atol=atol,
            dense_output=True,
        )

        if not res.success:
            raise RuntimeError(f"ODE integration failed: {res.message}")

        t = res.t
        n = res.y
        # Clip tiny numerical undershoots to enforce physical positivity.
        n = np.maximum(n, 0.0)
        x_liquid = np.zeros_like(n)
        C_gas = np.zeros_like(n)
        y_gas = np.zeros_like(n)
        OAV = np.zeros_like(n)

        for idx in range(n.shape[1]):
            n_col = n[:, idx]
            n_total = n_col.sum()
            if n_total > 1e-18:
                x = n_col / n_total
                # Clamp tiny numerical negatives to zero and re-normalize.
                x = np.clip(x, 0.0, None)
                x = x / x.sum()
            else:
                x = np.zeros_like(n_col)
            x_liquid[:, idx] = x
            C_gas[:, idx] = self.vle.gas_concentrations_mg_m3(x)
            y_gas[:, idx] = self.vle.gas_mole_fractions(x)
            OAV[:, idx] = self._oav(x, C_gas[:, idx])

        return Trajectory(
            t=t,
            n_liquid=n,
            x_liquid=x_liquid,
            C_gas=C_gas,
            y_gas=y_gas,
            OAV=OAV,
            formula=self.formula,
            notes={
                "method": method,
                "surface_area_m2": self.surface_area,
                "k_base_m_s": self.k_base,
                "k_skin_s-1": self._k_skin.tolist(),
                "n_eval": len(t_eval),
                "ode_message": res.message,
            },
        )

    def _oav(self, x: np.ndarray, C_gas: np.ndarray) -> np.ndarray:
        """Odor Activity Value = C_gas_i / odor_threshold_i.

        C_gas is in mg/m³.  odor_threshold_ug_m3 is in µg/m³.
        Convert C_gas to µg/m³ (×1000) before dividing so the units match.
        """
        thresholds = np.array(
            [i.odor_threshold_ug_m3 or np.nan for i in self.formula.ingredients],
            dtype=float,
        )
        with np.errstate(divide="ignore", invalid="ignore"):
            return np.where(thresholds > 0, (C_gas * 1000.0) / thresholds, 0.0)