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

# In[ ]:


from collections import OrderedDict
import warnings
from IPython.display import display
import sympy as sp
import matplotlib.pyplot as plt
from chempy.chemistry import Equilibrium, ReactionSystem, Substance
from chempy.kinetics.integrated import binary_rev
from chempy.kinetics.ode import get_odesys
get_ipython().run_line_magic('matplotlib', 'inline')
sp.init_printing()


# In[ ]:


t, kf, kb, X, Y, Z = sp.symbols('t k_f k_b X Y Z')
binary_rev(t, kf, kb, Y, Z, X, sp)


# In[ ]:


eq = Equilibrium.from_string('Fe+3 + SCN- = FeSCN+2; 10**2.065')
rsys = ReactionSystem(eq.as_reactions(kf=900))
rsys


# In[ ]:


odesys, extra = get_odesys(rsys)


# In[ ]:


def plot_against_analytic(odsy, c0, tend=1, integrator='cvode', atol=1e-12, rtol=1e-10):
    xout, yout, info = odsy.integrate(tend, c0, integrator=integrator, atol=atol, rtol=rtol)
    ref = binary_rev(xout[1:], rsys.rxns[0].param, rsys.rxns[1].param, c0['FeSCN+2'], c0['Fe+3'], c0['SCN-'])
    {k: v for k, v in info.items() if not k.startswith('internal')}
    plt.figure(figsize=(14, 4))
    plt.subplot(1, 2, 1)
    plt.plot(xout[1:], ref, alpha=0.3, linewidth=3)
    _ = odsy.plot_result()
    plt.subplot(1, 2, 2)
    plt.plot(xout[1:], yout[1:, odsy.names.index('FeSCN+2')] - ref)
    ref[:3], ref[-3:]
    return {k: v for k, v in info.items() if not k.startswith('internal')}


# In[ ]:


conc = {'Fe+3': 10e-3, 'SCN-': 2e-3, 'FeSCN+2': 2e-4}
plot_against_analytic(odesys, conc)


# Below are the derivations for ``binary_rev``:

# In[ ]:


x, a, b, c, U, V, S = sp.symbols('x a b c U V S')
u, v = 2*a*x + b, sp.sqrt(b**2 - 4*a*c)  # b**2 > 4*a*c
Prim = sp.log((U-V)/(U+V))/V  # primitive recip. 2nd order polynomial
prim = sp.log((u-v)/(u+v))/v  # primitive recip. 2nd order polynomial
prim.diff(x).simplify()


# In[ ]:


y = Y + X - x
z = Z + X - x


# In[ ]:


dxdt = kf*y*z - kb*x
dxdt


# In[ ]:


dxdt.expand()


# In[ ]:


expr2 = dxdt.expand().collect(x)
expr2


# In[ ]:


coeffs = expr2.as_poly(x).coeffs()
coeffs


# In[ ]:


eq = (sp.exp(Prim) - sp.exp(t + S))
eq


# In[ ]:


sp.Eq(Prim*V, (t+S)*V)


# In[ ]:


eq = sp.exp(Prim*V) - sp.exp((t+S)*V)
eq


# In[ ]:


eq2 = eq.subs({U: u})
eq2


# In[ ]:


sol, = sp.solve(eq2, x)
sol


# In[ ]:


s, = sp.solve(sol.subs(t, 0) - X, S)
s


# In[ ]:


sol2 = sol.subs(S, s).simplify()
sol2


# In[ ]:


R, W = sp.symbols('R W')
r = b + 2*X*a
w = sp.exp(-V*t)
sol3 = sol2.subs({w: W, r: R}).simplify().collect(W)
sol3


# In[ ]:


exprs = [
    sp.Eq(x, sol3),
    sp.Eq(R, r),
    sp.Eq(W, w),
    sp.Eq(V, v),
    sp.Eq(a, coeffs[0]),
    sp.Eq(b, coeffs[1]),
    sp.Eq(c, coeffs[2]),
]
for expr in exprs:
    display(expr)


# In[ ]:


master = exprs[0]
for e in exprs[1:]:
    master = master.subs(e.lhs, e.rhs)
master


# In[ ]:


cses, (reform,) = sp.cse(master)
reform


# In[ ]:


cses


# In[ ]:





# In[ ]:


ret = 'return ' + str(reform.rhs).replace('k_f', 'kf').replace('k_b', 'kb')
for key, val in cses:
    print('%s = %s' % (key, str(val).replace('k_f', 'kf').replace('k_b', 'kb').replace('exp', 'be.exp').replace('sqrt', 'be.sqrt')))
print(ret)


# In[ ]:


with warnings.catch_warnings():
    warnings.filterwarnings('error')
    binary_rev(1.0, 1e10, 1e-4, 0, 3e-7, 2e-7)  # check that our formulation avoids overflow


# Below we will test ``get_native``:

# In[ ]:


from chempy.kinetics._native import get_native


# In[ ]:


print({k: v.composition for k, v in rsys.substances.items()})
rsys.substances = OrderedDict([(k, Substance.from_formula(v.name)) for k, v in rsys.substances.items()])
print({k: v.composition for k, v in rsys.substances.items()})


# In[ ]:


nsys = get_native(rsys, odesys, 'gsl')


# In[ ]:


plot_against_analytic(nsys, conc, integrator='gsl', atol=1e-16, rtol=1e-16)


# In[ ]: