#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import matplotlib.pyplot as plt
import numpy as np
from chempy import ReactionSystem
from chempy.units import to_unitless, SI_base_registry as si, default_units as u, default_constants as const
from chempy.kinetics.ode import get_odesys
get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


R = 8.314472
T_K = 300
dH=80e3
dS=10
rsys1b = ReactionSystem.from_string("""
NO + Br -> NOBr; EyringParam(dH={dH}*J/mol, dS={dS}*J/K/mol)
""".format(dH=dH, dS=dS))
c0 = 1 # mol/dm3 === 1000 mol/m3
kbref = 20836643994.118652*T_K*np.exp(-(dH - T_K*dS)/(R*T_K))/c0
kbref


# In[ ]:


NO0_M = 1.5
Br0_M = 0.7
init_cond = dict(
    NOBr=0*u.M,
    NO=NO0_M*u.M,
    Br=Br0_M*u.M
)
t = 5*u.second
params = dict(
    temperature=T_K*u.K
)


# In[ ]:


def analytic_b(t):
    U, V = NO0_M, Br0_M
    d = U - V
    return (U*(1 - np.exp(-kbref*t*d)))/(U/V - np.exp(-kbref*t*d))


# In[ ]:


def integrate_and_plot(rsys):
    odesys, extra = get_odesys(rsys, unit_registry=si, constants=const)
    fig, axes = plt.subplots(1, 4, figsize=(14, 6))
    res = odesys.integrate(t, init_cond, params, integrator='cvode')
    t_sec = to_unitless(res.xout, u.second)
    NOBr_ref = analytic_b(t_sec)
    cmp = to_unitless(res.yout, u.M)
    ref = np.empty_like(cmp)
    ref[:, odesys.names.index('NOBr')] = NOBr_ref
    ref[:, odesys.names.index('Br')] = Br0_M - NOBr_ref
    ref[:, odesys.names.index('NO')] = NO0_M - NOBr_ref
    axes[0].plot(t_sec, cmp)
    axes[1].plot(t_sec, ref)
    axes[2].plot(t_sec, cmp - ref)
    res.plot_invariant_violations(ax=axes[3])
    assert np.allclose(cmp, ref)
    print({k: v for k, v in res.info.items() if not k.startswith('internal')})    


# In[ ]:


integrate_and_plot(rsys1b)


# In[ ]:


rsys2b = ReactionSystem.from_string("""
NO + Br -> NOBr; MassAction(EyringHS([{dH}*J/mol, {dS}*J/K/mol]))
""".format(dH=dH, dS=dS))


# In[ ]:


integrate_and_plot(rsys2b)


# In[ ]: