"""
Methanol synthesis loop — reactor, separator, recycle and purge
bdrworks.org · Handbook · Methanol, worked model 3

Make-up gas from the reformer is mixed with recycle gas and fed to the tube-cooled
reactor of model 2. The reactor outlet is cooled to 40 °C; methanol and water
condense (taken here as a perfect split), and the remaining gas is recycled.
A small purge keeps CH4 and N2 — which enter with the make-up gas and never
react — from accumulating without limit.

Unknowns: the recycle flows R_i (7 species).  Equations (tear stream):
    R = (1 - phi) * G( reactor( M + R ) )
where G keeps the gas species and phi is the purge fraction. Solved with fsolve.

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

# ---------------------------------------------------------------- data
SPECIES = ["CO", "CO2", "H2", "CH3OH", "H2O", "CH4", "N2"]
#                 CO   CO2   H2  MeOH  H2O  CH4  N2
NU = np.array([[0.0, -1.0, -3.0, 1.0, 1.0, 0.0, 0.0],     # 1: CO2 + 3H2 <=> CH3OH + H2O
               [1.0, -1.0, -1.0, 0.0, 1.0, 0.0, 0.0]])    # 2: CO2 + H2  <=> CO + H2O (RWGS)
RG = 8.314                                                 # J/(mol K)

# ideal-gas cp = A + B T + C T^2 + D T^3, J/(mol K)  (Reid, Prausnitz & Poling)
CP = np.array([[30.87, -1.285e-2, 2.789e-5, -1.272e-8],   # CO
               [19.80, 7.344e-2, -5.602e-5, 1.715e-8],    # CO2
               [27.14, 9.274e-3, -1.381e-5, 7.645e-9],    # H2
               [21.15, 7.092e-2, 2.587e-5, -2.852e-8],    # CH3OH
               [32.24, 1.924e-3, 1.055e-5, -3.596e-9],    # H2O
               [19.25, 5.213e-2, 1.197e-5, -1.132e-8],    # CH4
               [31.15, -1.357e-2, 2.680e-5, -1.168e-8]])  # N2
HF = np.array([-110.53, -393.51, 0.0, -200.94, -241.83, -74.87, 0.0]) * 1e3   # J/mol, 298.15 K


def cp(T):
    """Heat capacities of all species, J/(mol K)."""
    return CP @ np.array([1.0, T, T**2, T**3])


def h(T):
    """Molar enthalpies H_f(298) + integral of cp, J/mol."""
    T0 = 298.15
    dT = np.array([T - T0, (T**2 - T0**2) / 2, (T**3 - T0**3) / 3, (T**4 - T0**4) / 4])
    return HF + CP @ dT


def dH_rxn(T):
    """Heats of reaction 1 and 2, J/mol."""
    return NU @ h(T)


def K_eq(T):
    """Equilibrium constants (Graaf et al.), partial pressures in bar.
    K1 for CO2 + 3H2 = CH3OH + H2O (bar^-2);  K3 for CO + H2O = CO2 + H2 (-)."""
    K1 = 10.0 ** (3066.0 / T - 10.592)
    K3 = 10.0 ** (2073.0 / T - 2.029)
    return K1, K3


def rates(T, p):
    """Vanden Bussche & Froment (1996) rates, mol/(kg_cat s); p = partial pressures in bar."""
    pCO, pCO2, pH2, pM, pW = [max(x, 1e-10) for x in p[:5]]
    k = lambda A, B: A * np.exp(B / (RG * T))
    k1, k2, k3, k4, k5 = k(1.07, 36696), k(3453.38, 0), k(0.499, 17197), k(6.62e-11, 124119), k(1.22e10, -94765)
    K1, K3 = K_eq(T)
    den = 1.0 + k2 * pW / pH2 + k3 * np.sqrt(pH2) + k4 * pW
    r1 = k1 * pCO2 * pH2 * (1.0 - pW * pM / (K1 * pH2**3 * pCO2)) / den**3
    r2 = k5 * pCO2 * (1.0 - K3 * pW * pCO / (pCO2 * pH2)) / den
    return np.array([r1, r2])


def equilibrium(T, P, F):
    """Chemical equilibrium of reactions 1 and 2 from flows F (any unit); returns outlet flows."""
    from scipy.optimize import fsolve
    K1, K3 = K_eq(T)
    def res(x):
        n = F + NU.T @ x
        n = np.maximum(n, 1e-12)
        p = P * n / n.sum()
        return [np.log(p[3] * p[4] / (p[1] * p[2]**3)) - np.log(K1),
                np.log(p[0] * p[4] / (p[1] * p[2])) + np.log(K3)]
    x = fsolve(res, [0.1 * F.sum() * 0.05, 0.0], xtol=1e-12)
    return F + NU.T @ x

# ---------------------------------------------------------------- reactor (as in model 2)
P = 75.0; T_IN = 498.0; TC = 523.0
N_T = 5000; D = 0.040; L = 7.0; RHO_B = 1140.0; U = 250.0
A_T = np.pi * D**2 / 4


def rhs(z, u):
    F, T = u[:7], u[7]
    r = rates(T, P * F / F.sum())
    dF = RHO_B * A_T * (NU.T @ r)
    dT = (RHO_B * A_T * np.dot(-dH_rxn(T), r) - U * np.pi * D * (T - TC)) / np.dot(F, cp(T))
    return np.append(dF, dT)


def reactor(F_in):
    """F_in in kmol/h (whole reactor) -> outlet flows kmol/h and hot-spot temperature."""
    f = F_in * 1000 / 3600 / N_T
    sol = solve_ivp(rhs, (0, L), np.append(f, T_IN), method="LSODA", rtol=1e-7, atol=1e-10)
    return sol.y[:7, -1] * N_T * 3600 / 1000, sol.y[7].max()


# ---------------------------------------------------------------- loop
M_FLOW = 2500.0                                    # make-up gas, kmol/h
M_Y = np.array([0.17, 0.08, 0.72, 0.0, 0.0, 0.02, 0.01])
M = M_FLOW * M_Y
GAS = np.array([1, 1, 1, 0, 0, 1, 1], dtype=float) # species that stay in the gas at 40 °C


def solve_loop(phi, R_guess):
    def tear(R):
        out, _ = reactor(M + R)
        return (1 - phi) * GAS * out - R
    R = fsolve(tear, R_guess, xtol=1e-9)
    out, Tmax = reactor(M + R)
    return R, out, Tmax


def report(phi, R, out):
    meoh = out[3]                                   # kmol/h condensed methanol
    c_in = M[0] + M[1]
    return dict(recycle_ratio=R.sum() / M_FLOW, meoh=meoh, tpd=meoh * 32.04 * 24 / 1000,
                carbon_eff=meoh / c_in, inerts=(R[5] + R[6]) / R.sum(), purge=phi * (GAS * out).sum())


if __name__ == "__main__":
    SN = (M_Y[2] - M_Y[1]) / (M_Y[0] + M_Y[1])
    print(f"make-up {M_FLOW:.0f} kmol/h, stoichiometric number SN = {SN:.2f}")
    R = 3.0 * M_FLOW * np.array([0.05, 0.04, 0.70, 0, 0, 0.15, 0.06])
    phi = 0.05
    R, out, Tmax = solve_loop(phi, R)
    d = report(phi, R, out)
    print(f"purge {100*phi:.0f} %: recycle ratio {d['recycle_ratio']:.2f}, inerts in loop {100*d['inerts']:.1f} %, "
          f"hot spot {Tmax:.1f} K")
    print(f"methanol {d['meoh']:.0f} kmol/h = {d['tpd']:.0f} t/d, carbon efficiency {100*d['carbon_eff']:.1f} %")

    # purge fraction sweep (continuation from the previous solution)
    phis = [0.02, 0.03, 0.04, 0.05, 0.07, 0.10, 0.14, 0.20]
    rows, Rg = [], R
    for ph in sorted(phis, key=lambda x: abs(x - phi)):
        Rg, o, Tm = solve_loop(ph, Rg)
        rows.append((ph, report(ph, Rg, o), Tm))
    rows.sort(key=lambda r: r[0])
    for ph, dd, Tm in rows:
        print(f"  purge {100*ph:4.0f} %: recycle {dd['recycle_ratio']:.2f}, inerts {100*dd['inerts']:5.1f} %, "
              f"{dd['tpd']:5.0f} t/d, carbon eff. {100*dd['carbon_eff']:.1f} %")

    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    x = [100 * r[0] for r in rows]
    ax[0].plot(x, [100 * r[1]["carbon_eff"] for r in rows], "o-", lw=2, label="carbon efficiency (%)")
    ax[0].plot(x, [100 * r[1]["inerts"] for r in rows], "s-", lw=2, label="CH$_4$ + N$_2$ in loop (%)")
    ax[0].set(xlabel="purge fraction of separator gas (%)", title="Purge: efficiency vs inert build-up"); ax[0].legend()
    ax[1].plot(x, [r[1]["recycle_ratio"] for r in rows], "o-", lw=2, c="C2")
    ax[1].set(xlabel="purge fraction of separator gas (%)", ylabel="recycle / make-up (mol/mol)", title="Recycle ratio")
    for a in ax:
        a.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig("meoh_loop.png", dpi=150)
    plt.show()