import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
import matplotlib.pyplot as plt
from skopt import gp_minimize

# ============================================================
# CONSTANTS
# ============================================================
R = 8.314
R_gas = 0.08206
T_fixed = 308.15

V_liq = 0.015
V_gas = 0.005
kla_CO2 = 0.0092
kla_NH3 = 0.0092

# ============================================================
# PARAMETERS
# ============================================================
Ea = 35300
T_star = 414.6
KM = 3.21e-3
KES1 = 7.57e-7
KES2 = 1.27e-8
KP = 1.22e-2

parameter_kinetic = 40.04
urease = 25 * 0.015 / 1000 * parameter_kinetic  # FIXED (important)

# ============================================================
# THERMODYNAMICS
# ============================================================
def temperature_constants(T):
    KNH3 = np.exp(191.97 - 8451.61/T - 31.4335*np.log(T) + 0.0152123*T)
    KCO2 = np.exp(2767.92 - 80063.5/T - 478.653*np.log(T) + 0.714984*T)
    KHCO3 = np.exp(12.405 - 6286.89/T - 0.050628*T)
    KH2O = np.exp(14.01708 - 10294.83/T - 0.039282*T)
    H_CO2 = np.exp(-(1082.37 - 34417.2/T - 182.28*np.log(T) + 0.25159*T))
    H_NH3 = np.exp(-(160.559 - 8621.06/T - 25.6767*np.log(T) + 0.035388*T))
    return KNH3, KCO2, KHCO3, KH2O, H_CO2, H_NH3

KNH3, KCO2, KHCO3, KH2O, H_CO2, H_NH3 = temperature_constants(T_fixed)

# ============================================================
# pH MODEL
# ============================================================
def electroneutrality(pH, TA, TC):
    H = 10**(-pH)
    OH = KH2O / H
    NH4 = TA * (H * KNH3) / (KNH3 * H + KH2O)
    HCO3 = TC * (KCO2 * H) / (H**2 + KCO2 * H + KCO2 * KHCO3)
    CO3 = TC * (KCO2 * KHCO3) / (H**2 + KCO2 * H + KCO2 * KHCO3)
    return NH4 + H - (HCO3 + 2*CO3 + OH)

def solve_pH(TA, TC):
    TA = max(TA, 1e-16)
    TC = max(TC, 1e-16)
    return brentq(electroneutrality, 2, 12, args=(TA, TC))

# ============================================================
# MODEL
# ============================================================
def model(t, y):
    UREA, TA, TC, n_NH3, n_CO2 = y

    pH = solve_pH(TA, TC)
    H = 10**(-pH)

    P_CO2 = n_CO2 * R_gas * T_fixed / V_gas
    P_NH3 = n_NH3 * R_gas * T_fixed / V_gas

    CO2_aq = TC * H**2 / (H**2 + KCO2*H + KCO2*KHCO3)
    NH3_free = TA * KH2O / (KNH3*H + KH2O)

    rCO2 = kla_CO2 * (CO2_aq - H_CO2 * P_CO2)
    rNH3 = kla_NH3 * (NH3_free - H_NH3 * P_NH3)

    S = max(UREA, 1e-16)
    P = max(NH3_free, 1e-16)

    kT = np.exp(-Ea/R * (1/T_fixed - 1/T_star))

    ve = kT * urease * S / (
        (KM + S) *
        (1 + P / KP) *
        (1 + H / KES1 + KES2 / H)
    )

    return [
        -ve,
        2*ve - rNH3,
        ve - rCO2,
        rNH3 * V_liq,
        rCO2 * V_liq
    ]

# ============================================================
# SIMULATION
# ============================================================
def simulate(P_CO2_init, U0):
    n_CO2 = P_CO2_init * V_gas / (R_gas * T_fixed)

    y0 = [U0, 0.0, 0.0, 0.0, n_CO2]

    sol = solve_ivp(
        model,
        (0, 60),
        y0,
        method="BDF",
        t_eval=[60],
        rtol=1e-6,
        atol=1e-9
    )

    if not sol.success:
        return 1e6

    return sol.y[0, -1]

# ============================================================
# OPTIMIZATION WRAPPER
# ============================================================
def objective(x, U0):
    return simulate(x[0], U0)

# ============================================================
# LOOP OVER INITIAL UREA
# ============================================================
urea_init_range = np.linspace(0.015, 0.045, 10)

optimal_pressures = []
final_ureas = []

for U0 in urea_init_range:
    print(f"Optimizing U0 = {U0:.4f}")

    res = gp_minimize(
        lambda x: objective(x, U0),
        [(0, 3.5)],
        n_calls=50,
        random_state=42,
        acq_func="EI"
    )

    P_opt = res.x[0]

    # IMPORTANT: recompute true model value
    U_final = simulate(P_opt, U0)

    optimal_pressures.append(P_opt)
    final_ureas.append(U_final)

# ============================================================
# PLOT (YOUR STYLE)
# ============================================================
x = urea_init_range
values1 = np.array(optimal_pressures)
values2 = np.array(final_ureas)

plt.rcParams.update({
    "font.family": "serif",
    "font.size": 12,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.linewidth": 1,
})

fig, ax = plt.subplots(figsize=(5, 4))

# CO2 pressure
ax.plot(x, values1, 'o-', color='black')
ax.set_xlabel(r"$c_\mathrm{U}(0)$ / mol L$^{-1}$")
ax.set_ylabel(r"$p_{\mathrm{CO_2}}(0)$ / atm")
ax.set_xticks(np.linspace(0.015, 0.045, 4))
ax.set_yticks(np.linspace(0, 4, 6))

for spine in ax.spines.values():
    spine.set_visible(True)
    spine.set_linewidth(1.1)

# final urea
ax2 = ax.twinx()
ax2.plot(x, values2, 'o-', color='blue')
ax2.set_ylabel(r"$c_\mathrm{U}(60 min)$ / mol L$^{-1}$", color='blue')
ax2.tick_params(axis='y', colors='blue')

ax2.set_ylim(0, max(values2) * 1.1)
ax2.set_yticks(np.linspace(0, max(values2), 6))
ax2.set_ylim(0, 0.040 + 0)

# reasonable automatic limits
ax2.set_yticks(np.linspace(0, 0.04, 6))
plt.tight_layout()
plt.savefig("combined_results.pdf", bbox_inches="tight")
plt.show()