"""
Air separation — phase equilibrium of air and the pressures of a double column
bdrworks.org · Handbook · Air separation, worked model 1

Vapour pressures of N2, Ar and O2 from the Lee–Kesler corresponding-states correlation,
    ln Pr_sat = ln Pr0(Tr) + w ln Pr1(Tr),  w fitted to the normal boiling point,
and Raoult's law K_i = Psat_i/P for the ternary N2–Ar–O2 (CO2 and water are removed upstream).
Dew and bubble points of air follow from sum(y/K) = 1 and sum(K x) = 1.
Latent heats come from the Clapeyron equation, dH = T dV dPsat/dT, with the vapour and liquid
volumes of the Peng–Robinson equation (check against measured nitrogen data).
The double column couples the two pressures: nitrogen condensing in the HP column must be
dT (about 2 K) warmer than oxygen boiling in the LP column.

Run:  python air_equilibrium.py      (needs numpy, scipy, matplotlib)
"""
import numpy as np
from scipy.optimize import brentq
import matplotlib.pyplot as plt

RG = 8.314462618
SP = ["N2", "Ar", "O2"]
# Tc (K), Pc (bar), normal boiling point Tn (K), omega (for Peng–Robinson)
CRIT = {"N2": (126.2, 34.00, 77.35, 0.038), "Ar": (150.9, 48.98, 87.3, 0.000), "O2": (154.6, 50.43, 90.2, 0.022)}
AIR = np.array([0.78118, 0.00934, 0.20957]); AIR /= AIR.sum()        # dry, CO2-free air


def lnPr0(Tr):
    return 5.92714 - 6.09648/Tr - 1.28862*np.log(Tr) + 0.169347*Tr**6


def lnPr1(Tr):
    return 15.2518 - 15.6875/Tr - 13.4721*np.log(Tr) + 0.43577*Tr**6


# Lee–Kesler acentric factor that reproduces the normal boiling point
W_LK = {s: (np.log(1.01325/CRIT[s][1]) - lnPr0(CRIT[s][2]/CRIT[s][0]))/lnPr1(CRIT[s][2]/CRIT[s][0]) for s in SP}


def psat(s, T):
    Tc, Pc = CRIT[s][:2]
    Tr = T/Tc
    return Pc*np.exp(lnPr0(Tr) + W_LK[s]*lnPr1(Tr))                  # bar


def tsat(s, P):
    return brentq(lambda T: psat(s, T) - P, 50.0, CRIT[s][0] - 1e-6)


def K(T, P):
    return np.array([psat(s, T) for s in SP])/P


def dew_T(y, P):
    return brentq(lambda T: (y/K(T, P)).sum() - 1.0, 55.0, 150.0)


def bubble_T(x, P):
    return brentq(lambda T: K(T, P) @ x - 1.0, 55.0, 150.0)


def pr_volumes(s, T, P):
    """Liquid and vapour compressibility factors from Peng–Robinson (P in bar)."""
    Tc, Pc, _, w = CRIT[s]
    kap = 0.37464 + 1.54226*w - 0.26992*w*w
    a = 0.45724*(RG*Tc)**2/(Pc*1e5)*(1 + kap*(1 - np.sqrt(T/Tc)))**2
    b = 0.07780*RG*Tc/(Pc*1e5)
    A, B = a*P*1e5/(RG*T)**2, b*P*1e5/(RG*T)
    r = np.roots([1.0, -(1 - B), A - 3*B*B - 2*B, -(A*B - B*B - B**3)])
    r = np.sort(r[np.abs(r.imag) < 1e-10].real)
    return r[0], r[-1]


def latent_heat(s, T):
    """Clapeyron: dH = R T^2 (Zv - Zl) dln(Psat)/dT, J/mol."""
    h = 1e-3
    dlnP = (np.log(psat(s, T + h)) - np.log(psat(s, T - h)))/(2*h)
    zl, zv = pr_volumes(s, T, psat(s, T))
    return RG*T*T*(zv - zl)*dlnP


if __name__ == "__main__":
    print("Lee–Kesler omega:", {s: round(W_LK[s], 4) for s in SP})
    print("Nitrogen vapour pressure against measured data (Ullmann, Nitrogen, Table 1):")
    for T, P_meas in [(63.15, 0.1246), (80, 1.367), (100, 7.777), (120, 25.13)]:
        print(f"  {T:6.2f} K: {psat('N2', T):7.4f} bar vs {P_meas} bar ({100*(psat('N2', T)/P_meas - 1):+.2f} %)")
    for T, dH_meas in [(80, 2218.0 + 3248.0), (100, 2466.0 + 2033.0), (120, 2103.0 + 504.6)]:
        print(f"  latent heat {T} K: {latent_heat('N2', T):.0f} J/mol vs {dH_meas:.0f} J/mol")
    print("normal boiling points:", {s: round(tsat(s, 1.01325), 2) for s in SP})
    for P in [1.0, 6.0, 8.0]:
        print(f"air at {P:.0f} bar: dew point {dew_T(AIR, P):.2f} K, bubble point {bubble_T(AIR, P):.2f} K")

    # double column: O2 boils in the LP sump, N2 condenses at the HP top
    PL = 1.3
    TO2 = tsat("O2", PL)
    for dT in [1.0, 1.5, 2.0, 2.5]:
        print(f"LP {PL} bar (O2 boils at {TO2:.2f} K), dT = {dT} K -> HP column at {psat('N2', TO2 + dT):.2f} bar")
    print(f"latent heats in the main condenser: N2 at HP {latent_heat('N2', TO2 + 2):.0f} J/mol, O2 at LP {latent_heat('O2', TO2):.0f} J/mol")

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.4))
    Ps = np.linspace(1.0, 10.0, 60)
    ax[0].plot(Ps, [dew_T(AIR, P) for P in Ps], lw=2, label="dew point of air (starts to condense)")
    ax[0].plot(Ps, [bubble_T(AIR, P) for P in Ps], lw=2, label="bubble point of air (fully liquid)")
    ax[0].plot([1.0, 6.0], [81.5, 101.0], "ks", label="condensation temperatures quoted by Häussinger et al.")
    ax[0].set(xlabel="pressure (bar)", ylabel="temperature (K)", title="Liquefying air")
    ax[0].legend(fontsize=8); ax[0].grid(alpha=0.3)

    for pair, c in [(("N2", "O2"), "tab:blue"), (("Ar", "O2"), "tab:orange"), (("N2", "Ar"), "tab:green")]:
        al = [psat(pair[0], bubble_T(AIR, P))/psat(pair[1], bubble_T(AIR, P)) for P in Ps]
        ax[1].plot(Ps, al, lw=2, color=c, label=f"{pair[0]}/{pair[1]}")
    ax[1].set(xlabel="pressure (bar)", ylabel="relative volatility α (at the bubble point of air)", title="Volatilities fall with pressure")
    ax[1].legend(); ax[1].grid(alpha=0.3)

    Ts = np.linspace(70, 120, 200)
    ax[2].semilogy(Ts, [psat("N2", T) for T in Ts], lw=2, label="nitrogen")
    ax[2].semilogy(Ts, [psat("O2", T) for T in Ts], lw=2, label="oxygen")
    ax[2].semilogy(Ts, [psat("Ar", T) for T in Ts], lw=1.5, ls="--", label="argon")
    PH = psat("N2", TO2 + 2.0)
    ax[2].plot([TO2, TO2 + 2.0], [PL, PH], "k-", lw=1)
    ax[2].plot([TO2], [PL], "ko"); ax[2].plot([TO2 + 2.0], [PH], "ko")
    ax[2].annotate(f"LP column\nO₂ boils at {PL} bar", (TO2, PL), xytext=(95, 0.45), fontsize=8, arrowprops=dict(arrowstyle="-", lw=0.6))
    ax[2].annotate(f"HP column\nN₂ condenses at {PH:.1f} bar", (TO2 + 2.0, PH), xytext=(75, 9), fontsize=8, arrowprops=dict(arrowstyle="-", lw=0.6))
    ax[2].set(xlabel="temperature (K)", ylabel="vapour pressure (bar)", title="Coupling of the two columns (ΔT = 2 K)")
    ax[2].legend(loc="lower right"); ax[2].grid(alpha=0.3, which="both")
    fig.tight_layout(); fig.savefig("air_equilibrium.png", dpi=150)
    plt.show()