from __future__ import (absolute_import, division, print_function)
from functools import reduce, partial
from operator import mul
import sympy as sp
import numpy as np
import matplotlib.pyplot as plt
from pyneqsys.symbolic import SymbolicSys, TransformedSys, linear_exprs
sp.init_printing()

"""
Explanation: Row-reduce and variable transformations in non-linear equation systems, an applied example: Chemical equilibria
One of the strengths of pyneqsys is its ability to represent the system of non-linear equations symbolically. This allows the user to reformulate the problem in an physically equivalent but in a numerically different form.
In this notebook we will look at how we can remove linearly dependent equations automatically and go from an overdetermined system to a system with equal number of unknowns as equations. The latter is the preferred form (when it's possible to achive) since it gives a square Jacboian matrix and there are a larger family of numerial methods which we can use to optimize it (i.e. root finding).
End of explanation
"""


texnames = 'H^+ OH^- NH_4^+ NH_3 H_2O'.split()
n = len(texnames)
NH3_idx = texnames.index('NH_3')
NH3_varied = np.logspace(-7, 0)
c0 = 1e-7, 1e-7, 1e-7, 1, 55
K = Kw, Ka = 10**-14/55, 10**-9.24

"""
Explanation: Let's consider:
$$ \rm
H_2O \rightleftharpoons H^+ + OH^- \
NH_4^+ \rightleftharpoons H^+ + NH_3
$$
End of explanation
"""


stoichs = [[1, 1, 0, 0, -1], [1, 0, -1, 1, 0]]  # our 2 equilibria
H = [1, 1, 4, 3, 2]
N = [0, 0, 1, 1, 0]
O = [0, 1, 0, 0, 1]
q = [1, -1, 1, 0, 0]  # charge
preserv = [H, N, O, q]

"""
Explanation: Let's define the stoichiometry and composition:
End of explanation
"""


prod = lambda x: reduce(mul, x)

def get_f(x, params, backend, lnK):
    init_concs = params[:n]
    eq_constants = params[n:]
    le = linear_exprs(preserv, x, linear_exprs(preserv, init_concs), rref=True)
    if lnK:
        return le + [
            sum(backend.log(xi)*p for xi, p in zip(x, coeffs)) - backend.log(K) 
            for coeffs, K in zip(stoichs, eq_constants)
        ]
    else:
        return le + [
            prod(xi**p for xi, p in zip(x, coeffs)) - K for coeffs, K in zip(stoichs, eq_constants)
        ]

"""
Explanation: and now a function for the system of equations:
End of explanation
"""


neqsys = SymbolicSys.from_callback(
    partial(get_f, lnK=False), n, n+len(K),
    latex_names=[r'\mathrm{[%s]}' % nam for nam in texnames],
    latex_param_names=[r'\mathrm{[%s]_0}' % nam for nam in texnames] + [r'K_{\rm w}', r'K_{\rm a}(\mathrm{NH_4^+})']
)
neqsys

neqsys.get_jac()

%matplotlib inline
def solve_and_plot(nsys):
    fig = plt.figure(figsize=(12, 4))
    ax_out = plt.subplot(1, 2, 1, xscale='log', yscale='log')
    ax_err = plt.subplot(1, 2, 2, xscale='log')
    ax_err.set_yscale('symlog', linthreshy=1e-14)
    xres, extra = nsys.solve_and_plot_series(
        c0, c0+K, NH3_varied, NH3_idx, 'scipy', 
        plot_kwargs=dict(ax=ax_out), plot_residuals_kwargs=dict(ax=ax_err))
    for ax in (ax_out, ax_err):
        ax.set_xlabel('[NH3]0 / M')
    ax_out.set_ylabel('Concentration / M')
    ax_out.legend(loc='best')
    ax_err.set_ylabel('Residuals')
    
    avg_nfev = np.average([nfo['nfev'] for nfo in extra['info']])
    avg_njev = np.average([nfo['njev'] for nfo in extra['info']])
    success = np.average([int(nfo['success']) for nfo in extra['info']])
    return {'avg_nfev': avg_nfev, 'avg_njev': avg_njev, 'success': success}

    
solve_and_plot(neqsys)

"""
Explanation: note how we passed rref=True to linear_exprs, this will give a linear system in reduced row echolon form the system of equations. The four preservation equations (one for charge and three for atom types) has one linearly dependent equation which is dropped by pyneqsys.symbolic.linear_exprs, and after adding our two equations from the equilibria we are left with 5 equations (same number as unknowns).
End of explanation
"""


tneqsys = TransformedSys.from_callback(
    partial(get_f, lnK=True), (sp.exp, sp.log), 5, 7,
    latex_names=neqsys.latex_names, latex_param_names=neqsys.latex_param_names)
tneqsys

"""
Explanation: Now let's see how pyneqsys can transform our system:
End of explanation
"""


c_res, info = tneqsys.solve([1]*5, np.array(c0+K))
c0, c_res, info['success']

solve_and_plot(tneqsys)

"""
Explanation: Note how the conservation laws became non-linear while the expressions corresponding to the equilibria became linear.
End of explanation
"""


