"""
Peng-Robinson EOS - From-scratch implementation with kij.

Uses thermo library only for critical properties (Tc, Pc, omega) and
ideal gas Cp correlations. All EOS calculations (mixing rules, cubic
solver, fugacity coefficients, departure enthalpies) are from scratch
with binary interaction parameters from kij_peng_robinson.

Reference: Chapitre 5 - Equation d'etat de Peng-Robinson
"""

import os
import sys
import numpy as np

_src_dir = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
if _src_dir not in sys.path:
    sys.path.insert(0, _src_dir)

from thermo import ChemicalConstantsPackage
from kij_peng_robinson import build_kij_matrix, COMP_INDEX

R = 8.314  # J/(mol·K)
SQRT2 = np.sqrt(2.0)

_NAME_MAP = {
    "C1": "methane", "C2": "ethane", "C3": "propane",
    "iC4": "isobutane", "nC4": "butane",
    "iC5": "isopentane", "nC5": "pentane",
    "C6": "hexane", "C7": "heptane", "C8": "octane",
    "C9": "nonane", "C10": "decane", "C11": "undecane", "C12": "dodecane",
    "CO2": "carbon dioxide", "N2": "nitrogen",
}

COMPONENT_DB = {
    "C1": {"Mw": 16.04}, "C2": {"Mw": 30.07}, "C3": {"Mw": 44.10},
    "iC4": {"Mw": 58.12}, "nC4": {"Mw": 58.12},
    "iC5": {"Mw": 72.15}, "nC5": {"Mw": 72.15},
    "C6": {"Mw": 86.18}, "C7": {"Mw": 100.21}, "C8": {"Mw": 114.23},
    "C9": {"Mw": 128.25}, "C10": {"Mw": 142.28}, "C11": {"Mw": 156.31},
    "C12": {"Mw": 170.34}, "CO2": {"Mw": 44.01}, "N2": {"Mw": 28.01},
}


class PengRobinsonEOS:
    """From-scratch PR EOS with binary interaction parameters (kij)."""

    def __init__(self, component_names):
        self.names = list(component_names)
        self.nc = len(self.names)

        # Critical properties, Cp_ig, and Hvap from thermo database
        thermo_names = [_NAME_MAP[n] for n in self.names]
        constants, props = ChemicalConstantsPackage.from_IDs(thermo_names)
        self.Tc = np.array(constants.Tcs)
        self.Pc = np.array(constants.Pcs)
        self.omega = np.array(constants.omegas)
        self.Mw = np.array([COMPONENT_DB[n]["Mw"] for n in self.names])
        self._Cp_ig = props.HeatCapacityGases
        self._Hvap = props.EnthalpyVaporizations

        # kij submatrix for our components
        full_kij = build_kij_matrix()
        self.kij = np.zeros((self.nc, self.nc))
        for i, ni in enumerate(self.names):
            for j, nj in enumerate(self.names):
                self.kij[i, j] = full_kij[COMP_INDEX[ni], COMP_INDEX[nj]]

        # PR pure-component constants
        self.ai0 = 0.45724 * R**2 * self.Tc**2 / self.Pc
        self.bi = 0.07780 * R * self.Tc / self.Pc
        self.kappa = 0.37464 + 1.54226 * self.omega - 0.26992 * self.omega**2

    # ================================================================
    # Alpha function and derivative
    # ================================================================
    def _alpha(self, T):
        """alpha_i(T) = [1 + kappa*(1 - sqrt(T/Tc))]^2"""
        f = 1.0 + self.kappa * (1.0 - np.sqrt(max(T, 1.0) / self.Tc))
        return f * f

    def _dalpha_dT(self, T):
        """d(alpha_i)/dT"""
        T = max(T, 1.0)
        sqrt_Tr = np.sqrt(T / self.Tc)
        f = 1.0 + self.kappa * (1.0 - sqrt_Tr)
        return -self.kappa * f / np.sqrt(T * self.Tc)

    # ================================================================
    # Mixing rules with kij
    # ================================================================
    def _mix_params(self, z, T):
        """
        Compute a_m, b_m, da_m/dT with van der Waals one-fluid mixing rules.
        Returns (am, bm, dam_dT, ai_T, aij).
        """
        z = np.asarray(z, dtype=float)
        z = np.maximum(z, 1e-30)
        z = z / np.sum(z)

        alpha = self._alpha(T)
        ai_T = self.ai0 * alpha
        sqrt_ai = np.sqrt(np.maximum(ai_T, 0.0))

        # a_ij = sqrt(ai*aj) * (1-kij)
        aij = np.outer(sqrt_ai, sqrt_ai) * (1.0 - self.kij)

        am = z @ aij @ z
        bm = z @ self.bi

        # da_m/dT for enthalpy departure
        dalpha = self._dalpha_dT(T)
        dai_T = self.ai0 * dalpha

        nc = self.nc
        daij = np.zeros((nc, nc))
        for i in range(nc):
            for j in range(nc):
                sij = sqrt_ai[i] * sqrt_ai[j]
                if sij > 1e-30:
                    daij[i, j] = (1.0 - self.kij[i, j]) * \
                        (dai_T[i] * ai_T[j] + ai_T[i] * dai_T[j]) / (2.0 * sij)
        dam_dT = z @ daij @ z

        return am, bm, dam_dT, ai_T, aij

    # ================================================================
    # Cubic solver
    # ================================================================
    def _solve_cubic(self, A, B):
        """
        Solve PR cubic: Z^3 - (1-B)Z^2 + (A-3B^2-2B)Z - (AB-B^2-B^3) = 0.
        Returns (Z_V, Z_L).

        When only 1 real root exists (supercritical mixture), the real
        part of the complex conjugate pair is used as the "virtual"
        root for the missing phase.
        """
        c2 = -(1.0 - B)
        c1 = A - 3.0 * B * B - 2.0 * B
        c0 = -(A * B - B * B - B * B * B)

        roots = np.roots([1.0, c2, c1, c0])

        real_roots = []
        complex_real_part = None
        for r in roots:
            if abs(r.imag) < 1e-8 and r.real > 0:
                real_roots.append(r.real)
            elif complex_real_part is None and r.real > 0:
                complex_real_part = r.real

        if len(real_roots) >= 2:
            real_roots.sort()
            Z_V = real_roots[-1]
            Z_L = max(real_roots[0], B + 1e-6)
            return Z_V, Z_L

        if len(real_roots) == 1:
            Z_real = real_roots[0]
            if complex_real_part is not None and complex_real_part > B:
                Z_virtual = max(complex_real_part, B + 1e-6)
                if Z_real > Z_virtual:
                    return Z_real, Z_virtual
                else:
                    return Z_virtual, Z_real
            return Z_real, max(B + 1e-6, Z_real * 0.15)

        Z = max(B + 0.01, 0.3)
        return Z, Z

    # ================================================================
    # Fugacity coefficients
    # ================================================================
    def fugacity_coefficients(self, T, P, z, phase='vapor'):
        """
        Compute ln(phi_i) for each component.

        ln(phi_i) = (bi/bm)(Z-1) - ln(Z-B)
                   - A/(2*sqrt(2)*B) * [2*sum_j(zj*aij)/am - bi/bm]
                     * ln[(Z+(1+sqrt2)*B)/(Z+(1-sqrt2)*B)]
        """
        am, bm, _, ai_T, aij = self._mix_params(z, T)

        A = am * P / (R * T) ** 2
        B = bm * P / (R * T)

        Z_V, Z_L = self._solve_cubic(A, B)
        Z = Z_V if phase == 'vapor' else Z_L

        z = np.maximum(np.asarray(z, dtype=float), 1e-30)
        z = z / np.sum(z)

        sum_za = aij @ z  # sum_j z_j*a_ij for each i

        zmB = max(Z - B, 1e-30)
        arg = max(Z + (1.0 + SQRT2) * B, 1e-30) / max(Z + (1.0 - SQRT2) * B, 1e-30)
        ln_ratio = np.log(max(arg, 1e-30))

        fac = A / (2.0 * SQRT2 * max(B, 1e-30))
        am_safe = max(am, 1e-30)
        bm_safe = max(bm, 1e-30)

        ln_phi = np.zeros(self.nc)
        for i in range(self.nc):
            bi_bm = self.bi[i] / bm_safe
            ln_phi[i] = (bi_bm * (Z - 1.0)
                         - np.log(zmB)
                         - fac * (2.0 * sum_za[i] / am_safe - bi_bm) * ln_ratio)

        return ln_phi

    def K_values(self, T, P, x, y):
        """K_i = phi_L_i / phi_V_i from PR fugacity coefficients."""
        ln_phi_L = self.fugacity_coefficients(T, P, x, 'liquid')
        ln_phi_V = self.fugacity_coefficients(T, P, y, 'vapor')
        return np.exp(np.clip(ln_phi_L - ln_phi_V, -50, 50))

    def K_values_wilson(self, T, P):
        """Wilson correlation for initial K estimates."""
        return (self.Pc / P) * np.exp(5.373 * (1.0 + self.omega) * (1.0 - self.Tc / T))

    # ================================================================
    # Enthalpy
    # ================================================================
    def enthalpy_departure(self, T, P, z, phase='vapor'):
        """
        PR departure enthalpy (J/mol).

        H_dep = RT(Z-1) + (T*da/dT - a)/(2*sqrt(2)*b)
                * ln[(Z+(1+sqrt2)*B)/(Z+(1-sqrt2)*B)]
        """
        am, bm, dam_dT, _, _ = self._mix_params(z, T)

        A = am * P / (R * T) ** 2
        B = bm * P / (R * T)

        Z_V, Z_L = self._solve_cubic(A, B)
        Z = Z_V if phase == 'vapor' else Z_L

        arg = max(Z + (1.0 + SQRT2) * B, 1e-30) / max(Z + (1.0 - SQRT2) * B, 1e-30)
        ln_ratio = np.log(max(arg, 1e-30))

        H_dep = R * T * (Z - 1.0)
        if bm > 1e-30:
            H_dep += (T * dam_dT - am) / (2.0 * SQRT2 * bm) * ln_ratio

        return H_dep

    def ideal_gas_enthalpy(self, T, z):
        """Ideal gas enthalpy (J/mol) from Cp_ig integration, T_ref=298.15 K."""
        z = np.maximum(np.asarray(z, dtype=float), 1e-30)
        z = z / np.sum(z)
        T_ref = 298.15
        H_ig = 0.0
        for i in range(self.nc):
            if z[i] > 1e-20:
                try:
                    H_ig += z[i] * self._Cp_ig[i].T_dependent_property_integral(T_ref, T)
                except Exception:
                    H_ig += z[i] * 35.0 * (T - T_ref)  # rough fallback
        return H_ig

    def mixture_enthalpy(self, T, P, z, phase='vapor'):
        """Total mixture enthalpy (J/mol) = H_ig + H_dep (PR EOS)."""
        return self.ideal_gas_enthalpy(T, z) + self.enthalpy_departure(T, P, z, phase)

    def pure_component_enthalpy(self, T, P, z, phase='vapor'):
        """
        Mixture enthalpy from pure component enthalpies (ideal mixing).

        This is the approach used by ProSim, Mixprops, and similar tools:
        - Vapor: H = sum_i z_i * H_ig_i(T)
        - Liquid: h = sum_i z_i * [H_ig_i(T) - dHvap_i(T)]

        For supercritical components (T >= Tc_i), dHvap = 0 so h_i = H_ig_i.
        Cp_ig and dHvap correlations from DIPPR (via thermo library).
        """
        z = np.maximum(np.asarray(z, dtype=float), 1e-30)
        z = z / np.sum(z)
        T_ref = 298.15

        H = 0.0
        for i in range(self.nc):
            if z[i] < 1e-20:
                continue

            # Ideal gas enthalpy: integral(Cp_ig, T_ref, T)
            try:
                H_ig_i = self._Cp_ig[i].T_dependent_property_integral(T_ref, T)
            except Exception:
                H_ig_i = 35.0 * (T - T_ref)

            if phase == 'liquid' and T < self.Tc[i]:
                # Subcritical: subtract latent heat of vaporization
                try:
                    dHvap = self._Hvap[i].T_dependent_property(T)
                    if dHvap is not None and dHvap > 0:
                        H += z[i] * (H_ig_i - dHvap)
                    else:
                        H += z[i] * H_ig_i
                except Exception:
                    H += z[i] * H_ig_i
            else:
                # Vapor, or supercritical liquid component
                H += z[i] * H_ig_i

        return H

    # ================================================================
    # Flash TP (Rachford-Rice + successive substitution)
    # ================================================================
    def flash_TP(self, T, P, z):
        """TP flash. Returns (VF, x, y, K)."""
        z = np.maximum(np.asarray(z, dtype=float), 1e-30)
        z = z / np.sum(z)

        K = self.K_values_wilson(T, P)

        for outer in range(30):
            if np.sum(z * K) <= 1.0:
                return 0.0, z.copy(), z.copy(), K
            if np.sum(z / K) <= 1.0:
                return 1.0, z.copy(), z.copy(), K

            VF = 0.5
            for _ in range(100):
                Km1 = K - 1.0
                d = 1.0 + VF * Km1
                f = np.sum(z * Km1 / d)
                df = -np.sum(z * Km1 ** 2 / d ** 2)
                if abs(df) < 1e-30:
                    break
                dV = -f / df
                VF = np.clip(VF + dV, 1e-10, 1.0 - 1e-10)
                if abs(dV) < 1e-12:
                    break

            d = 1.0 + VF * (K - 1.0)
            x = z / d
            x /= np.sum(x)
            y = K * x
            y /= np.sum(y)

            K_new = self.K_values(T, P, x, y)
            if np.max(np.abs(np.log(np.maximum(K_new, 1e-30) / np.maximum(K, 1e-30)))) < 1e-6:
                K = K_new
                break
            K = 0.5 * K + 0.5 * K_new

        d = 1.0 + VF * (K - 1.0)
        x = z / d
        x /= np.sum(x)
        y = K * x
        y /= np.sum(y)
        return VF, x, y, K

    # ================================================================
    # Bubble point temperature
    # ================================================================
    def bubble_T(self, P, x, T_guess=None):
        """Bubble point T by Newton-Raphson with PR K-values."""
        x = np.maximum(np.asarray(x, dtype=float), 1e-30)
        x = x / np.sum(x)

        T = float(T_guess) if T_guess else np.dot(x, self.Tc) * 0.7
        y = x.copy()

        for _ in range(60):
            K = self.K_values(T, P, x, y)
            y = K * x
            ys = np.sum(y)
            y = y / max(ys, 1e-30)
            f = ys - 1.0

            dT = max(0.1, abs(T) * 1e-4)
            K_p = self.K_values(T + dT, P, x, y)
            f_p = np.sum(K_p * x) - 1.0
            df = (f_p - f) / dT
            if abs(df) < 1e-30:
                break
            delta = np.clip(-f / df, -30, 30)
            T = np.clip(T + delta, 100, 700)
            if abs(delta) < 0.01:
                break

        K = self.K_values(T, P, x, y)
        y = K * x
        y /= max(np.sum(y), 1e-30)
        return T, K, y
