"""
Air separation — membrane nitrogen generator and the energy of the three routes
bdrworks.org · Handbook · Air separation, worked model 3

Compressed air flows along hollow fibres; O2 permeates faster than N2 and the N2-rich
retentate is the product. Binary O2/N2 model (argon counted with nitrogen), permeate at
1 bar. Local permeate composition (Weller–Steiner):
    y/(1-y) = alpha (x - phi y) / ((1 - x) - phi (1 - y)),     phi = p_permeate/p_feed
Complete mixing: the whole feed side is at the retentate composition.
Cross flow: each element's permeate leaves without mixing; along the fibre
    d(n)/d(x) = n/(y - x),     dA = -dn / J,   J = Q_O2 (p_f x - p_p y) + Q_N2 (p_f (1-x) - p_p (1-y))
The energy is the air compression, w = R T0 ln(p_f/p0) / eta_T per mol of air, and the
thermodynamic minimum follows from the ideal work of separating ideal gases.

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

RG, T0, VN = 8.314462618, 298.15, 0.022414        # J/(mol K), K, m3(STP)/mol
XF = 0.2095                                        # O2 in dry air
GPU = 3.3464e-10                                   # mol/(m2 s Pa) per GPU

# Table 7 of Häussinger et al. (Ullmann, Nitrogen): air at 12.35 bar
TABLE7 = [(0.005, 21.7, 0.73), (0.01, 27.3, 0.58), (0.02, 35.0, 0.45), (0.04, 46.1, 0.34), (0.06, 55.2, 0.29), (0.08, 63.6, 0.26)]
PF, PP = 12.35, 1.013


def y_local(x, a, phi):
    """Root in (0, 1) of phi(1 - a) y^2 + [1 + (a - 1)(x + phi)] y - a x = 0."""
    A, B, C = phi*(1 - a), 1 + (a - 1)*(x + phi), -a*x
    return -2*C/(B + np.sqrt(B*B - 4*A*C))


def cross_flow(a, phi, xr, xf=XF):
    """Retentate flow per mol of feed when the O2 fraction is brought from xf to xr."""
    sol = solve_ivp(lambda x, n: [n[0]/(y_local(x, a, phi) - x)], (xf, xr), [1.0], rtol=1e-10, atol=1e-12)
    return sol.y[0, -1]


def complete_mixing(a, phi, xr, xf=XF):
    y = y_local(xr, a, phi)
    theta = (xf - xr)/(y - xr)                    # stage cut
    return 1.0 - theta if 0 <= theta <= 1 else np.nan


def recovery(nr, xr, xf=XF):
    return nr*(1 - xr)/(1 - xf)


def cross_flow_area(a, QO2_gpu, pf, pp, xr, xf=XF):
    """Membrane area (m2) per mol/s of feed air, cross flow."""
    QO2, QN2 = QO2_gpu*GPU, QO2_gpu*GPU/a
    phi = pp/pf
    def rhs(x, u):
        n = u[0]; y = y_local(x, a, phi)
        dn = n/(y - x)
        J = QO2*(pf*1e5*x - pp*1e5*y) + QN2*(pf*1e5*(1 - x) - pp*1e5*(1 - y))
        return [dn, -dn/J]
    sol = solve_ivp(rhs, (xf, xr), [1.0, 0.0], rtol=1e-9, atol=1e-12)
    return sol.y[1, -1]


def w_min(yp, Pp, R, P0=1.0, z=(XF, 1 - XF)):
    """Ideal work (J per mol of product) to make product (O2 fraction yp) at Pp from air at P0,
    with nitrogen recovery R and the reject at P0 (ideal gases at T0)."""
    z = np.array(z); npr = R*z[1]/(1 - yp); nw = 1 - npr
    yw = (z[0] - npr*yp)/nw
    s = lambda y, P: sum(v*np.log(v*P/P0) for v in y if v > 0)
    return RG*T0*(npr*s([yp, 1 - yp], Pp) + nw*s([yw, 1 - yw], P0) - s(z, P0))/npr


def kwh_per_nm3(J_per_mol):
    return J_per_mol/VN/3.6e6


if __name__ == "__main__":
    phi = PP/PF
    fit = minimize_scalar(lambda a: sum((100*recovery(cross_flow(a, phi, xr), xr) - rec)**2 for xr, rec, _ in TABLE7),
                          bounds=(2.5, 8.0), method="bounded")
    a_fit = fit.x
    print(f"cross flow fitted to Table 7: alpha = {a_fit:.2f} (rms {np.sqrt(fit.fun/len(TABLE7)):.1f} points of recovery)")
    print("O2 %   Table 7   cross(fit)   cross(alpha 6)   complete(alpha 6)")
    for xr, rec, E in TABLE7:
        print(f"{100*xr:4.1f}   {rec:5.1f}    {100*recovery(cross_flow(a_fit, phi, xr), xr):6.1f}      "
              f"{100*recovery(cross_flow(6.0, phi, xr), xr):6.1f}        {100*recovery(complete_mixing(6.0, phi, xr), xr):6.1f}")
    x_lim = XF/(6.0 - 5.0*XF)
    print(f"complete mixing, alpha 6, phi -> 0: the retentate cannot fall below {100*x_lim:.2f} % O2")

    # energy: isothermal compression implied by the tables
    w_air = RG*T0*np.log(PF/1.013)
    etas = []
    for xr, rec, E in TABLE7:
        air_per_prod = (1 - xr)/(rec/100*(1 - XF))
        eta = kwh_per_nm3(w_air*air_per_prod)/E; etas.append(eta)
    print(f"Table 7 energies imply an isothermal compression efficiency of {min(etas):.2f}–{max(etas):.2f}")
    # cryogenic single column (Table 5): 0.25 kWh/m3 N2 at 8 bar; recovery from air_column.py (40 stages, 1 ppm): 64.0 %
    rec_col = 0.640; w_col = RG*T0*np.log(8.5/1.013)/(rec_col*(1 - XF))
    print(f"single column, 64 % recovery, air to 8.5 bar: isothermal work {kwh_per_nm3(w_col):.3f} kWh/Nm3 -> "
          f"Table 5 (0.25) implies eta = {kwh_per_nm3(w_col)/0.25:.2f}")
    for yp, Pp, R, lab in [(1e-6, 8.0, 0.64, "single column, 1 ppm, 8 bar, 64 %"), (1e-6, 8.0, 1e-6, "same, vanishing recovery"),
                           (0.01, 11.35, 0.273, "membrane, 1 % O2, 27 %"), (0.01, 11.35, 1e-6, "membrane product, vanishing recovery")]:
        print(f"minimum work, {lab}: {kwh_per_nm3(w_min(yp, Pp, R)):.4f} kWh/Nm3")
    A = cross_flow_area(6.0, 40.0, PF, PP, 0.01)
    nr = cross_flow(6.0, phi, 0.01); prod = 1000/3600/VN                   # 1000 Nm3/h of product, mol/s
    area = A*prod/nr
    print(f"alpha 6, 40 GPU O2 fibres (2 Barrer, 50 nm skin): {area:.0f} m2 for 1000 Nm3/h at 1 % O2 "
          f"= {area/(np.pi*500e-6*2.0)/1e6:.2f} million fibres of 0.5 mm x 2 m")

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.4))
    xs = np.linspace(0.004, 0.10, 60)
    ax[0].plot(100*xs, [100*recovery(cross_flow(6.0, phi, x), x) for x in xs], lw=2, label="cross flow, α = 6 (1990 fibre)")
    ax[0].plot(100*xs, [100*recovery(cross_flow(a_fit, phi, x), x) for x in xs], lw=2, label=f"cross flow, α = {a_fit:.1f} (fitted)")
    rc = [100*recovery(complete_mixing(6.0, phi, x), x) for x in xs]
    ax[0].plot(100*xs, [r if r > 0 else np.nan for r in rc], lw=2, ls="--", label="complete mixing, α = 6")
    ax[0].plot([100*t[0] for t in TABLE7], [t[1] for t in TABLE7], "ks", label="Table 7 (12.35 bar)")
    ax[0].set(xlabel="O₂ left in product nitrogen (%)", ylabel="N₂ recovery (%)", title="Membrane: purity against recovery", ylim=(0, 90))
    ax[0].legend(fontsize=8); ax[0].grid(alpha=0.3)

    eta = float(np.mean(etas))
    ax[1].loglog(100*xs, [kwh_per_nm3(w_air*(1 - x)/(recovery(cross_flow(a_fit, phi, x), x)*(1 - XF)))/eta for x in xs], lw=2,
                 label=f"membrane model, η = {eta:.2f}")
    ax[1].loglog([100*t[0] for t in TABLE7], [t[2] for t in TABLE7], "ks", label="membrane, Table 7")
    ax[1].loglog([3.0, 0.1], [0.2*1.2505, 0.5*1.2505], "D", color="tab:orange", ms=7, label="PSA (Sec. 4.2)")
    ax[1].loglog([1e-4, 1e-4], [0.25, 0.22], "o", color="tab:green", ms=7, label="cryogenic, Table 5 (1 ppm)")
    yps = np.logspace(-6, -0.9, 50)
    ax[1].loglog(100*yps, [kwh_per_nm3(w_min(y, 8.0, 1e-6)) for y in yps], "k--", lw=1.2, label="minimum work, product at 8 bar")
    ax[1].set(xlabel="O₂ in product nitrogen (%)", ylabel="specific energy (kWh per Nm³ N₂)", title="Energy of the three routes", ylim=(0.03, 1.5))
    ax[1].legend(fontsize=7, loc="center left", bbox_to_anchor=(0.0, 0.36)); ax[1].grid(alpha=0.3, which="both")

    xa = np.linspace(0.005, 0.08, 30)
    areas = [cross_flow_area(6.0, 40.0, PF, PP, x)*prod/cross_flow(6.0, phi, x) for x in xa]
    ax[2].semilogy(100*xa, areas, lw=2, color="tab:purple")
    ax[2].set(xlabel="O₂ in product nitrogen (%)", ylabel="membrane area (m²)", title="Area for 1000 Nm³/h, α = 6, 40 GPU")
    ax[2].grid(alpha=0.3, which="both")
    fig.tight_layout(); fig.savefig("air_membrane.png", dpi=150)
    plt.show()