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

# In[ ]:


import sympy as sp
import matplotlib.pyplot as plt
from chempy import ReactionSystem
from chempy.kinetics.ode import get_odesys
sp.init_printing()
get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


get_ipython().run_line_magic('pinfo', 'ReactionSystem')


# In[ ]:


rsys = ReactionSystem.from_string("""
-> H; 'p'
2 H -> H2; 'k2'
""", checks=('substance_keys', 'duplicate', 'duplicate_names'))
rsys


# In[ ]:


H, k, p, c1, t, H0  = sp.symbols('H k p c1 t H0', real=True, positive=True)


# In[ ]:


dHdt = p - 2*k*H**2
dHdt


# In[ ]:


analytic = sp.sqrt(p/k/2)*sp.tanh((c1 + t)*sp.sqrt(2*k*p))
analytic


# In[ ]:


(analytic.diff(t) - dHdt.subs(H, analytic)).simplify()


# In[ ]:


sol1, sol2 = sp.solve(analytic.subs(t, 0) - H0, c1)


# In[ ]:


res1 = analytic.subs(c1, sol1).simplify()
res1


# In[ ]:


res2 = analytic.subs(c1, sol2).simplify()
res2


# In[ ]:


def get_python_code(expr, be='np'):
    cses, new_expr = sp.cse(expr)
    snippet = '\n'.join(['%s = %s' % cse for cse in cses] + ['return %s' % new_expr[0]]).replace(
        'sqrt', '%s.sqrt' % be).replace(
        'log', '%s.log' % be).replace(
        'tanh', '%s.tanh' % be).replace(
        'exp', '%s.exp' % be)
    return snippet


# In[ ]:


print(get_python_code(res2))


# In[ ]:


print(get_python_code(res2.rewrite('exp'), ))


# In[ ]:


f1 = sp.lambdify([t, p, k, H0], res1)
f2 = sp.lambdify([t, p, k, H0], res2)


# In[ ]:


odesys, extra = get_odesys(rsys, include_params=False)


# In[ ]:


c0 = {'H': 42e-6, 'H2': 17e-6}
params = {'p': 314/3600*.998*2*1.036426856e-7, 'k2': 53/60}
result = odesys.integrate(7*3600, c0, params,
                          integrator='cvode')
result.plot()
plt.plot(result.xout, f2(result.xout, params['p'], params['k2'], c0['H']), ls=':')
plt.gca().set_ylim([0, 300e-6])


# In[ ]:


result.params


# In[ ]: