"""
Air separation — single-column nitrogen generator (cryogenic)
bdrworks.org · Handbook · Air separation, worked model 2

Purified air (saturated vapour, N2–Ar–O2) enters below the bottom stage of a column at 8 bar.
Nitrogen vapour leaves the top stage; the product D is taken as gas and the rest is condensed
against the evaporating bottom liquid and returned as reflux. The O2-rich bottom liquid is
throttled and evaporated in that condenser (residual gas, later expanded in a turbine).

Equilibrium-stage model with constant molar overflow (V = F = 1, L = 1 - D on every stage),
Raoult's law K-values from Lee–Kesler vapour pressures, solved with the Wang–Henke
bubble-point method: for fixed stage temperatures the component balances form a
tridiagonal system (Thomas algorithm); new temperatures are the bubble points.

Run:  python air_column.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"]
CRIT = {"N2": (126.2, 34.00, 77.35), "Ar": (150.9, 48.98, 87.3), "O2": (154.6, 50.43, 90.2)}
AIR = np.array([0.78118, 0.00934, 0.20957]); AIR /= AIR.sum()


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


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]
    return Pc*np.exp(lnPr0(T/Tc) + W_LK[s]*lnPr1(T/Tc))


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


def dlnpsat(s, T):
    """d ln Psat / dT of the Lee–Kesler correlation, 1/K."""
    Tc = CRIT[s][0]; Tr = T/Tc
    d0 = 6.09648/Tr**2 - 1.28862/Tr + 6*0.169347*Tr**5
    d1 = 15.6875/Tr**2 - 13.4721/Tr + 6*0.43577*Tr**5
    return (d0 + W_LK[s]*d1)/Tc


def Kmat(T, P):
    """K-values for a vector of stage temperatures: array (stages x 3)."""
    return np.column_stack([psat(s, T) for s in SP])/P


def bubble_T(X, P, T0):
    """Bubble points of all stages at once: Newton's method on ln(sum K x) = 0."""
    T = np.array(T0, dtype=float)
    for _ in range(60):
        Km = Kmat(T, P); s = (Km*X).sum(axis=1)
        dlog = (Km*X*np.column_stack([dlnpsat(sp, T) for sp in SP])).sum(axis=1)/s
        dT = np.clip(-np.log(s)/dlog, -10.0, 10.0)
        T = T + dT
        if np.abs(dT).max() < 1e-11:
            break
    return T


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


def thomas(a, b, c, d):
    """Solve a tridiagonal system (a: sub-, b: main, c: super-diagonal)."""
    n = len(d); cp = np.zeros(n); dp = np.zeros(n)
    cp[0], dp[0] = c[0]/b[0], d[0]/b[0]
    for i in range(1, n):
        m = b[i] - a[i]*cp[i - 1]
        cp[i] = c[i]/m if i < n - 1 else 0.0
        dp[i] = (d[i] - a[i]*dp[i - 1])/m
    x = np.zeros(n); x[-1] = dp[-1]
    for i in range(n - 2, -1, -1):
        x[i] = dp[i] - cp[i]*x[i + 1]
    return x


def column(N, D, P, z=AIR, tol=1e-9, itmax=3000):
    """N equilibrium stages (1 = top), vapour feed F = 1 below stage N, gaseous product D.
    Returns stage temperatures T, liquid x and vapour y (N x 3)."""
    L = 1.0 - D
    T = np.full(N, dew_T(z, P)); w, dT_old = 1.0, np.inf
    for it in range(itmax):
        Kj = Kmat(T, P)
        X = np.zeros((N, 3))
        for i in range(3):
            a = np.full(N, L); a[0] = 0.0
            b = -(L + Kj[:, i]); b[0] = -(L + D*Kj[0, i])    # top: reflux = condensed top vapour
            c = np.zeros(N); c[:-1] = Kj[1:, i]
            d = np.zeros(N); d[-1] = -z[i]
            X[:, i] = thomas(a, b, c, d)
        X = np.maximum(X, 1e-300)
        X /= X.sum(axis=1, keepdims=True)
        Tn = bubble_T(X, P, T)
        dT = np.abs(Tn - T).max()
        if dT < tol:
            T = Tn
            break
        if it > 0 and dT > dT_old:                      # oscillating near the pinch: damp the update
            w = max(0.5*w, 0.05)
        T, dT_old = T + w*(Tn - T), dT
    else:
        raise RuntimeError(f"no convergence for N = {N}, D = {D}")
    Y = Kmat(T, P)*X
    return T, X, Y


def pinch_D(P, z=AIR):
    """Largest product with no O2 in it: bottom liquid in equilibrium with the air feed."""
    Td = dew_T(z, P)
    xs = z/K(Td, P); xs /= xs.sum()
    return 1.0 - z[2]/xs[2], xs


def D_for_ppm(N, P, ppm=1.0):
    """Largest product flow that still meets the O2 specification (capped just below the pinch)."""
    f = lambda D: np.log(column(N, D, P)[2][0, 2]*1e6/ppm)
    Dhi = pinch_D(P)[0] - 0.002
    return Dhi if f(Dhi) < 0 else brentq(f, 0.20, Dhi, xtol=1e-6)


if __name__ == "__main__":
    P = 8.0
    Dmax, xs = pinch_D(P)
    print(f"pinch: at most D = {Dmax:.3f} mol N2 per mol air ({100*Dmax/AIR[0]:.1f} % recovery); "
          f"bottom liquid then O2 {100*xs[2]:.1f} %, N2 {100*xs[0]:.1f} %")
    rows = []
    for N in [20, 30, 40, 50, 60]:
        try:
            D = D_for_ppm(N, P)
        except ValueError:
            continue
        T, X, Y = column(N, D, P)
        rows.append((N, D, 100*D*Y[0, 0]/AIR[0], Y[0, 1]*1e6, 100*X[-1, 2]))
        print(f"N = {N}: 1 ppm O2 at D = {D:.4f}, N2 recovery {rows[-1][2]:.1f} %, Ar in product {rows[-1][3]:.0f} ppm, "
              f"bottom liquid {rows[-1][4]:.1f} % O2, T top {T[0]:.2f} K, bottom {T[-1]:.2f} K")

    # the condenser: crude liquid evaporates at a lower pressure, 2 K below condensing nitrogen
    N, D = 40, D_for_ppm(40, P)
    T, X, Y = column(N, D, P)
    Tcond = dew_T(Y[0], P)
    Pev = brentq(lambda p: dew_T(X[-1], p) - (Tcond - 2.0), 1.0, P)
    print(f"design N = 40, D = {D:.4f}: N2 condenses at {Tcond:.2f} K; the crude liquid "
          f"({100*X[-1, 2]:.1f} % O2) must evaporate at or below {Pev:.2f} bar")

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.4))
    for Nn, col in [(20, "tab:red"), (30, "tab:orange"), (40, "tab:green"), (60, "tab:blue")]:
        Ds = np.linspace(0.30, Dmax - 0.002, 25)
        ppm = [column(Nn, d, P)[2][0, 2]*1e6 for d in Ds]
        ax[0].semilogy(100*Ds/AIR[0], ppm, lw=2, color=col, label=f"{Nn} stages")
    ax[0].axvline(100*Dmax/AIR[0], color="k", ls=":", lw=1)
    ax[0].axhline(1.0, color="0.5", lw=0.8)
    ax[0].text(100*Dmax/AIR[0] - 0.5, 1e-4, "pinch at the\nbottom stage", ha="right", fontsize=8)
    ax[0].set(xlabel="N₂ recovery (% of the N₂ in the air)", ylabel="O₂ in product nitrogen (ppm)", title="Purity against recovery, 8 bar", ylim=(1e-6, 1e5))
    ax[0].legend(fontsize=8); ax[0].grid(alpha=0.3, which="both")

    st = np.arange(1, N + 1)
    ax[1].semilogx(X[:, 2], st, lw=2, label="O₂ in liquid")
    ax[1].semilogx(X[:, 1], st, lw=2, label="Ar in liquid")
    ax[1].semilogx(1 - X[:, 0], st, lw=1, ls="--", color="0.4", label="all but N₂")
    ax[1].invert_yaxis()
    ax[1].set(xlabel="mole fraction", ylabel="stage (1 = top)", title=f"Profiles, {N} stages, D = {D:.3f}")
    ax[1].legend(fontsize=8); ax[1].grid(alpha=0.3, which="both")

    Ns = [r[0] for r in rows]
    ax[2].plot(Ns, [r[2] for r in rows], "o-", lw=2, color="tab:green", label="N₂ recovery at 1 ppm O₂")
    ax[2].set_xlabel("equilibrium stages"); ax[2].set_ylabel("N₂ recovery (%)", color="tab:green")
    b = ax[2].twinx(); b.semilogy(Ns, [r[3] for r in rows], "s--", lw=1.5, color="tab:purple")
    b.set_ylabel("argon in product (ppm)", color="tab:purple")
    ax[2].set_title("More stages: more nitrogen, and more argon in it"); ax[2].grid(alpha=0.3)
    fig.tight_layout(); fig.savefig("air_column.png", dpi=150)
    plt.show()