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

# In[ ]:


from __future__ import absolute_import, division, print_function
from collections import defaultdict
from ipywidgets import interact
import matplotlib.pyplot as plt
from chempy import Reaction, Substance, ReactionSystem
from chempy.kinetics.ode import get_odesys
from chempy.kinetics.analysis import plot_reaction_contributions
from chempy.printing.tables import UnimolecularTable, BimolecularTable
from chempy.util.graph import rsys2graph
import sympy
sympy.init_printing()
get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


A, B, C, D = map(Substance, 'ABCD')
One = sympy.S.One
reactions = r0, r1, r2 = [
    Reaction({'A'}, {'B'}, 4*One/100, name='R1: A cons.'),
    Reaction({'B', 'C'}, {'A', 'C'}, 10**(4*One), name='R2: A reform.'),
    Reaction({'B': 2}, {'B', 'C'}, 3*10**(7*One), name='R3: C form.')
]
rsys = ReactionSystem(reactions, (A, B, C))
rsys


# In[ ]:


uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)
bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)
assert not (not_uni & not_bi), "There are only uni- & bi-molecular reactions in this set"


# In[ ]:


uni


# In[ ]:


bi


# In[ ]:


rsys2graph(rsys, 'robertson.png', save='.')
from IPython.display import Image; Image('robertson.png')


# In[ ]:


odesys, extra = get_odesys(rsys, include_params=True)
odesys.exprs


# In[ ]:


odesys.get_jac()


# In[ ]:


c0 = defaultdict(float, {'A': 1})
result = odesys.integrate(1e10, c0, integrator='cvode', nsteps=2000)
{k: v for k, v in result.info.items() if not k.startswith('internal')}


# In[ ]:


extra['rate_exprs_cb'](result.xout, result.yout)


# In[ ]:


result.plot(xscale='log', yscale='log')


# In[ ]:


fig, axes = plt.subplots(2, 2, figsize=(14, 6))
plot_reaction_contributions(result, rsys, extra['rate_exprs_cb'], 'AB', axes=axes[0, :])
plot_reaction_contributions(result, rsys, extra['rate_exprs_cb'], 'AB', axes=axes[1, :],
                            relative=True, yscale='linear')
plt.tight_layout()


# We could also have parsed the reactions from a string:

# In[ ]:


str_massaction = """
A -> B; 'k1'
B + C -> A + C; 'k2'
2 B -> B + C; 'k3'
"""


# In[ ]:


rsys3 = ReactionSystem.from_string(str_massaction, substance_factory=lambda formula: Substance(formula))


# In[ ]:


rsys3.substance_names()


# In[ ]:


odesys3, extra3 = get_odesys(rsys3, include_params=False, lower_bounds=[0, 0, 0])
extra3['param_keys'], extra3['unique']


# In[ ]:


odesys3.exprs, odesys3.params, odesys3.names, odesys3.param_names


# In[ ]:


def integrate_and_plot(A0=1.0, B0=0.0, C0=0.0, lg_k1=-2, lg_k2=4, lg_k3=7, lg_tend=9):
    plt.figure(figsize=(14, 4))
    tout, yout, info = odesys3.integrate(
        10**lg_tend, {'A': A0, 'B': B0, 'C': C0},
        {'k1': 10**lg_k1, 'k2': 10**lg_k2, 'k3': 10**lg_k3},
        integrator='cvode', nsteps=3000)
    plt.subplot(1, 2, 1)
    odesys3.plot_result(xscale='log', yscale='log')
    plt.legend(loc='best')
    plt.subplot(1, 2, 2)
    plt.plot(tout[tout<.05], yout[tout<.05, odesys3.names.index('B')])
    _ = plt.legend('best')
interact(integrate_and_plot) #, **kw)


# In[ ]:


# We could also have used SymPy to construct symbolic rates:
import sympy
rsys_sym = ReactionSystem.from_string("""
A -> B; sp.Symbol('k1')
B + C -> A + C; sp.Symbol('k2')
2 B -> B + C; sp.Symbol('k3')
""", rxn_parse_kwargs=dict(globals_={'sp': sympy}), substance_factory=lambda formula: Substance(formula))
odesys_sym, _ = get_odesys(rsys_sym, params=True)
for attr in 'exprs params names param_names'.split():
    print(getattr(odesys_sym, attr))


# For larger systems it is easy to loose track of what substances are actually playing a part, here the html tables can help (note the yellow background color):

# In[ ]:


rsys.substances['D'] = D
uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)
uni


# In[ ]:


bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)
bi


# In[ ]: