from collections import defaultdict
from itertools import chain
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from ipywidgets import interact
from chempy import Substance, Reaction, ReactionSystem
from chempy.kinetics.rates import Arrhenius, MassAction
from chempy.kinetics.ode import get_odesys
from chempy.printing.numbers import number_to_scientific_latex
from chempy.units import SI_base_registry, default_units as u
sp.init_printing()
%matplotlib inline

"""
Explanation: Parameter variation studies of kinetic systems
This notebook shows how one can explore the impact of a certain parameter on a kinetic model. We will also use units explicitly for our parameters.
End of explanation
"""


A, B, C = map(Substance, 'ABC')
r1 = Reaction({'A'}, {'B'}, MassAction(Arrhenius(unique_keys=('A1', 'Ea_R_1'))))
r2 = Reaction({'B'}, {'C'}, MassAction(Arrhenius(unique_keys=('A2', 'Ea_R_2'))))
rsys = ReactionSystem([r1, r2])
rsys

"""
Explanation: We will use a generic model representing a decay-chain with two decays:
End of explanation
"""


params = {'A1': 1e11/u.s, 'A2': 2e11/u.s, 'Ea_R_1': 8e3*u.K, 'Ea_R_2': 8.5e3*u.K, 'temperature': 300*u.K}
c0 = defaultdict(lambda: 0*u.molar, {'A': 1*u.molar})
variables = c0.copy()
variables.update(params)
rsys.rates(variables)

odesys, extra = get_odesys(rsys, include_params=False, lower_bounds=0)
print(dict(zip(odesys.dep, odesys.names)))
print(dict(zip(odesys.params, odesys.param_names)))
odesys.exprs

"""
Explanation: "Arrhenius" behaviour means that the rate of reaction depends exponentially on the inverse absolute temperature.
We will use units on all our parameters in this notebook. This will prevent us from incorrect conversions or using parameters of the wrong dimensionality where they don't belong:
End of explanation
"""


params2 = params.copy()
pk = 'Ea_R_2'
params2[pk] = [8.1e3, 8.2e3, 8.3e3]*u.K

"""
Explanation: Let's look at the influence of Ea_R_2, we will choose three temperatures: 8100, 8200 and 8300 K (this is all fictive so never mind the very high temperatures):
End of explanation
"""


res2 = odesys.integrate(7*u.s, c0, params2, integrator='cvode')
fig, axes = plt.subplots(1, len(res2), figsize=(14, 4))
for r, ax in zip(res2, axes):
    r.plot(ax=ax)
    ax.set_title('$%s = %s$' % (pk.replace('_', '\\_'), number_to_scientific_latex(r.named_param('Ea_R_2'))))

"""
Explanation: Running the integartion & plotting the result:
End of explanation
"""


def integrate_and_plot(T_C=25):
    res = odesys.integrate(7*u.s, c0, dict(params, temperature=(T_C+273.15)*u.K), integrator='cvode')
    res.plot()

interact(integrate_and_plot)

"""
Explanation: We can also use ipywidgets to get interactive controls:
End of explanation
"""


