from __future__ import annotations

import logging
from typing import Any

import numpy as np
from thermo.unifac import DOUFIP2006, DOUFSG, UNIFAC

from ..models import Formula
from .unifac import DOUFSG_AVAILABLE_GROUPS, apply_structural_fallback, validate_subgroups

logger = logging.getLogger("pino.thermo.vle")


class VLECalculator:
    """Compute non-ideal VLE via Dortmund-Modified UNIFAC with dynamic fallback."""

    def __init__(self, formula: Formula) -> None:
        self.formula = formula
        self._cache: dict[tuple[float, ...], np.ndarray] = {}
        self._has_unifac = False
        self.unifac: UNIFAC | None = None
        self._setup_unifac()

    def _setup_unifac(self) -> None:
        """Build a UNIFAC object from the ingredient DDB UNIFAC groups."""
        groups: list[dict[int, int]] = []
        for ing in self.formula.ingredients:
            raw = ing.unifac_groups or {}
            # Coerce string/int keys to integers for validation.
            int_groups: dict[int, int] = {}
            for gid, count in raw.items():
                try:
                    key = int(gid)
                except (ValueError, TypeError):
                    logger.warning("Skipping non-integer UNIFAC group id %s", gid)
                    continue
                int_groups[key] = int(count)

            valid, missing = validate_subgroups(int_groups)
            if not valid:
                cleaned = apply_structural_fallback(int_groups, missing)
                if not cleaned:
                    logger.warning(
                        "UNIFAC unavailable for %s (CAS %s); activity coefficients default to 1.0",
                        ing.name, ing.cas
                    )
                    groups.append({})
                    continue
                int_groups = {int(k): int(v) for k, v in cleaned.items()}

            groups.append(int_groups)

        # Require every ingredient to have at least one valid UNIFAC group.
        self._has_unifac = all(bool(g) for g in groups) and len(groups) == len(self.formula.ingredients)
        if not self._has_unifac:
            logger.debug("UNIFAC unavailable for formula; activity coefficients default to 1.0")
            return

        try:
            self.unifac = UNIFAC.from_subgroups(
                T=self.formula.temperature_k,
                xs=[1.0 / len(self.formula.ingredients)] * len(self.formula.ingredients),
                chemgroups=groups,
                version=1,
                interaction_data=DOUFIP2006,
                subgroups=DOUFSG,
            )
        except Exception as exc:
            logger.warning("UNIFAC setup failed (%s); falling back to ideal gammas=1", exc)
            self._has_unifac = False
            self.unifac = None

    def activity_coefficients(self, x: np.ndarray) -> np.ndarray:
        """Return gamma_i for liquid mole fractions x using UNIFAC."""
        key = tuple(np.round(x, 12).tolist())
        if key in self._cache:
            return self._cache[key]

        if not self._has_unifac or self.unifac is None:
            gammas = np.ones_like(x, dtype=float)
        else:
            try:
                new_unifac = self.unifac.to_T_xs(
                    T=self.formula.temperature_k,
                    xs=x.tolist(),
                )
                gammas = np.array(new_unifac.gammas(), dtype=float)
            except Exception as exc:
                logger.warning("UNIFAC gamma calculation failed (%s); using gamma=1", exc)
                gammas = np.ones_like(x, dtype=float)

        self._cache[key] = gammas
        return gammas

    def partial_pressures(self, x: np.ndarray) -> np.ndarray:
        """P_i = gamma_i * x_i * P_i^sat (non-ideal Raoult)."""
        gamma = self.activity_coefficients(x)
        p_sat = np.array([i.vapor_pressure_pa for i in self.formula.ingredients], dtype=float)
        return gamma * x * p_sat

    def gas_mole_fractions(self, x: np.ndarray) -> np.ndarray:
        """y_i = P_i / P_total (assuming gas phase ideal at ambient pressure)."""
        p_partial = self.partial_pressures(x)
        p_total = p_partial.sum()
        if p_total <= 0:
            return np.zeros_like(x)
        return p_partial / p_total

    def gas_concentrations_mg_m3(self, x: np.ndarray) -> np.ndarray:
        """Convert gas mole fractions to mass concentrations at ambient T and P."""
        y = self.gas_mole_fractions(x)
        # total molar concentration (mol/m3) = P / (R T)
        R = 8.314462618  # J/(mol K)
        c_total = self.formula.ambient_pressure_pa / (R * self.formula.temperature_k)  # mol/m3
        mws = np.array([i.molecular_weight for i in self.formula.ingredients], dtype=float)  # g/mol
        # mg/m3 = y_i * C_total * MW_i * 1000 mg/g
        return y * c_total * mws * 1000.0

    def fugacity_ratios(self, x: np.ndarray) -> np.ndarray:
        """Driving force for evaporation: f_i^L / (P * phi_i) ≈ gamma_i * x_i * P_i^sat / P."""
        return self.partial_pressures(x) / self.formula.ambient_pressure_pa