"""
Methanol synthesis — adiabatic multi-bed quench converter (ICI type)
bdrworks.org · Handbook · Methanol, worked model 1

Four adiabatic catalyst beds. Only part of the feed is preheated and sent to bed 1;
the rest stays cold (quench gas) and is injected between the beds so that every
bed starts at the same inlet temperature.
    dF_i/dW = sum_j nu_ij r_j               (W = catalyst mass, kg)
    (sum F_i cp_i) dT/dW = sum_j (-dH_j) r_j
The quench flow before each bed follows from an enthalpy balance of the mixing
point, and the split of the feed is found so that all of it is used.

Run:  python meoh_quench.py      (needs numpy, scipy, matplotlib)
"""
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
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

# ---------------------------------------------------------------- converter
P = 75.0                                   # bar
F0 = 10000.0                               # total feed, kmol/h
Y0 = np.array([0.10, 0.05, 0.70, 0.005, 0.001, 0.12, 0.024]); Y0 = Y0 / Y0.sum()
BEDS = [8e3, 10e3, 14e3, 18e3]             # catalyst per bed, kg (50 t in total)
T_BED = 503.0                              # inlet temperature of every bed, K
T_Q = 400.0                                # quench gas temperature, K


def rhs(W, u):
    F, T = u[:7], u[7]
    r = rates(T, P * F / F.sum()) * 3.6            # kmol/(kg h)
    dF = NU.T @ r
    dT = np.dot(-dH_rxn(T), r) * 1e3 / np.dot(F * 1e3, cp(T))
    return np.append(dF, dT)


def quench_needed(F, T):
    """Quench flow (kmol/h of feed at T_Q) that brings the mixture to T_BED."""
    Fq = F0 * Y0
    H_hot = np.dot(F, h(T) - h(T_BED))              # heat to remove, kJ/h * 1e-3 scale cancels
    H_cold = np.dot(Fq / F0, h(T_BED) - h(T_Q))     # heat taken up per kmol/h of quench
    return H_hot / H_cold


def run(f1, keep=False):
    """Run all beds with a fraction f1 of the feed going to bed 1."""
    F, used, prof, Wst = f1 * F0 * Y0, f1, [], 0.0
    for k, Wb in enumerate(BEDS):
        if k > 0:
            q = quench_needed(F, T_out)
            F = F + q * Y0; used += q / F0
        sol = solve_ivp(rhs, (0, Wb), np.append(F, T_BED), method="LSODA", rtol=1e-8, atol=1e-8, dense_output=True)
        if keep:
            w = np.linspace(0, Wb, 200); uu = sol.sol(w)
            prof.append((Wst + w, uu))
        F, T_out, Wst = sol.y[:7, -1], sol.y[7, -1], Wst + Wb
    return used, F, prof


if __name__ == "__main__":
    f1 = brentq(lambda f: run(f)[0] - 1.0, 0.05, 0.95, xtol=1e-6)
    used, F, prof = run(f1, keep=True)
    print(f"feed to bed 1: {100*f1:.1f} % of {F0:.0f} kmol/h at {T_BED:.0f} K; quench gas at {T_Q:.0f} K")
    for k, (w, uu) in enumerate(prof, 1):
        y = uu[3] / uu[:7].sum(axis=0)
        ne = equilibrium(uu[7, -1], P, uu[:7, 0])
        print(f"bed {k}: {BEDS[k-1]/1e3:4.0f} t, T {uu[7,0]:.0f} -> {uu[7,-1]:.1f} K, MeOH out {100*y[-1]:.2f} % "
              f"(equilibrium at T_out {100*ne[3]/ne.sum():.2f} %)")
    meoh = F[3] - F0 * Y0[3]
    xc = 1 - (F[0] + F[1]) / (F0 * (Y0[0] + Y0[1]))
    print(f"production {meoh:.0f} kmol/h = {meoh*32.04*24/1000:.0f} t/d, CO+CO2 conversion {100*xc:.1f} %")

    # methanol - temperature diagram with the equilibrium line of the bed-1 feed
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    for k, (w, uu) in enumerate(prof, 1):
        ax[0].plot(uu[7], 100 * uu[3] / uu[:7].sum(axis=0), lw=2, label=f"bed {k}")
        ax[1].plot(w / 1e3, uu[7], lw=2, label=f"bed {k}")
    Te = np.linspace(480, 620, 40)
    ye = [100 * (lambda n: n[3] / n.sum())(equilibrium(t, P, f1 * F0 * Y0)) for t in Te]
    ax[0].plot(Te, ye, "--", c="gray", label="equilibrium (bed-1 feed)")
    ax[0].set(xlabel="temperature (K)", ylabel="CH$_3$OH (mol-%)", title="Adiabatic beds with cold-shot quench")
    ax[1].set(xlabel="cumulative catalyst (t)", ylabel="temperature (K)", title="Temperature profile")
    for a in ax:
        a.grid(alpha=0.3); a.legend()
    fig.tight_layout()
    fig.savefig("meoh_quench.png", dpi=150)
    plt.show()
