from operator import mul
from functools import reduce
from itertools import product
import chempy
from chempy.chemistry import Species, Equilibrium
from chempy.equilibria import EqSystem, NumSysLog, NumSysLin
import numpy as np
import sympy as sp
sp.init_printing()
import matplotlib.pyplot as plt
%matplotlib inline
sp.__version__, chempy.__version__

"""
Explanation: Two coupled equilibria: protolysis of ammonia in water
In this notebook we will look at how ChemPy can be used to formulate a system of (non-linear) equations from conservation laws and equilibrium equations. We will look att ammonia since it is a fairly well-known subtance.
End of explanation
"""


substance_names = ['H+', 'OH-', 'NH4+', 'NH3', 'H2O']
subst = {n: Species.from_formula(n) for n in substance_names}
assert [subst[n].charge for n in substance_names] == [1, -1, 1, 0, 0], "Charges of substances"
print(u'Composition of %s: %s' % (subst['NH3'].unicode_name, subst['NH3'].composition))

"""
Explanation: First, let's define our substances. ChemPy can parse chemical formulae into chempy.Substance instances with information on charge, composition, molar mass (and pretty-printing):
End of explanation
"""


init_conc = {'H+': 1e-7, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}
x0 = [init_conc[k] for k in substance_names]
H2O_c = init_conc['H2O']
w_autop = Equilibrium({'H2O': 1}, {'H+': 1, 'OH-': 1}, 10**-14/H2O_c)
NH4p_pr = Equilibrium({'NH4+': 1}, {'H+': 1, 'NH3': 1}, 10**-9.26)
equilibria = w_autop, NH4p_pr
[(k, init_conc[k]) for k in substance_names]

eqsys = EqSystem(equilibria, subst)
eqsys

"""
Explanation: Let's define som e initial concentrations and governing equilibria:
End of explanation
"""


x, sol, sane = eqsys.root(init_conc)
x, sol['success'], sane

"""
Explanation: We can solve the non-linear system of equations by calling the root method (the underlying representation uses pyneqsys:
End of explanation
"""


logx, logsol, sane = eqsys.root(init_conc, NumSys=(NumSysLog,))
logx, logsol['success'], sane

"""
Explanation: This system is quite easy to solve for, if we have convergence problems we can try to solve a transformed system. As an example we will use NumSysLog:
End of explanation
"""


x - logx

"""
Explanation: In thise case they give essentially the same answer:
End of explanation
"""


ny = len(substance_names)
y = sp.symarray('y', ny)
i = sp.symarray('i', ny)
K = Kw, Ka = sp.symbols('K_w K_a')
w_autop.param = Kw
NH4p_pr.param = Ka
ss = sp.symarray('s', ny)
ms = sp.symarray('m', ny)

numsys_log = NumSysLog(eqsys, backend=sp)
f = numsys_log.f(y, list(i)+list(K))
f

numsys_lin = NumSysLin(eqsys, backend=sp)
numsys_lin.f(y, i)

A, ks = eqsys.stoichs_constants(False, backend=sp)
[reduce(mul, [b**e for b, e in zip(y, row)]) for row in A]

from pyneqsys.symbolic import SymbolicSys
subs = list(zip(i, x0)) + [(Kw, 10**-14), (Ka, 10**-9.26)]
numf = [_.subs(subs) for _ in f]
neqs = SymbolicSys(list(y), numf)
neqs.solve([0, 0, 0, 0, 0], solver='scipy')

j = sp.Matrix(1, len(f), lambda _, q: f[q]).jacobian(y)
init_conc_j = {'H+': 1e-10, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}
xj = eqsys.as_per_substance_array(init_conc_j)
jarr = np.array(j.subs(dict(zip(y, xj))).subs({Kw: 1e-14, Ka: 10**-9.26}).subs(
            dict(zip(i, xj))))
jarr = np.asarray(jarr, dtype=np.float64)
np.log10(np.linalg.cond(jarr))

j.simplify()
j

eqsys.composition_balance_vectors()

numsys_rref_log = NumSysLog(eqsys, True, True, backend=sp)
numsys_rref_log.f(y, list(i)+list(K))

np.set_printoptions(4, linewidth=120)
scaling = 1e8
for rxn in eqsys.rxns:
    rxn.param = rxn.param.subs({Kw: 1e-14, Ka: 10**-9.26})

x, res, sane = eqsys.root(init_conc, rref_equil=True, rref_preserv=True)
x, res['success'], sane

x, res, sane = eqsys.root(init_conc, x0=eqsys.as_per_substance_array(
        {'H+': 1e-11, 'OH-': 1e-3, 'NH4+': 1e-3, 'NH3': 1.0, 'H2O': 55.5}))
res['success'], sane

x, res, sane = eqsys.root(init_conc, x0=eqsys.as_per_substance_array(
        {'H+': 1.7e-11, 'OH-': 3e-2, 'NH4+': 3e-2, 'NH3': 0.97, 'H2O': 55.5}))
res['success'], sane

init_conc

nc=60
Hp_0 = np.logspace(-3, 0, nc)
def plot_rref(**kwargs):
    fig, axes = plt.subplots(2, 2, figsize=(16, 6), subplot_kw=dict(xscale='log', yscale='log'))
    return [eqsys.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': axes.flat[i]}, rref_equil=e,
                        rref_preserv=p, **kwargs) for i, (e, p) in enumerate(product(*[[False, True]]*2))]

res_lin = plot_rref(method='lm')
[all(_[2]) for _ in res_lin]

for col_id in range(len(substance_names)):
    for i in range(1, 4):
        plt.subplot(1, 3, i, xscale='log')
        plt.gca().set_yscale('symlog', linthreshy=1e-14)
        plt.plot(Hp_0, res_lin[0][0][:, col_id] - res_lin[i][0][:, col_id])
plt.tight_layout()

eqsys.plot_errors(res_lin[0][0], init_conc, Hp_0, 'H+')

init_conc, eqsys.ns

res_log = plot_rref(NumSys=NumSysLog)

eqsys.plot_errors(res_log[0][0], init_conc, Hp_0, 'H+')

res_log_lin = plot_rref(NumSys=(NumSysLog, NumSysLin))
eqsys.plot_errors(res_log_lin[0][0], init_conc, Hp_0, 'H+')

from chempy.equilibria import NumSysSquare
res_log_sq = plot_rref(NumSys=(NumSysLog, NumSysSquare))
eqsys.plot_errors(res_log_sq[0][0], init_conc, Hp_0, 'H+')

res_sq = plot_rref(NumSys=(NumSysSquare,), method='lm')
eqsys.plot_errors(res_sq[0][0], init_conc, Hp_0, 'H+')

x, res, sane = eqsys.root(x0, NumSys=NumSysLog, rref_equil=True, rref_preserv=True)
x, res['success'], sane

x, res, sane = eqsys.root(x, NumSys=NumSysLin, rref_equil=True, rref_preserv=True)
x, res['success'], sane

"""
Explanation: We can create symbolic representations of these systems of equations:
End of explanation
"""


