%matplotlib inline
import sympy as sp
sp.init_printing()
import numpy as np
import matplotlib.pyplot as plt
from pyneqsys.symbolic import SymbolicSys

"""
Explanation: Parameter variation
In this notebook we will look at how pyneqsys can solve a system of non-linear equations for different values of a parameter (pyneqsys will use the preceeding solution as starting guess in the consecutive solves).
End of explanation
"""


def f(x, params):
    eq = lambda y, fy: 1/(y+x['a'])**params['p'] + x['b'] - fy
    return [eq(0, 1), eq(1, 0)]
neqsys = SymbolicSys.from_callback(f, 2, 1, names='a b'.split(), param_names='p', x_by_name=True, par_by_name=True)
neqsys.exprs

x0, params = dict(a=0.5, b=-0.5), dict(p=1)
ab, info = neqsys.solve(x0, params)
ab, info

print('Residuals: %s'  % f(dict(zip(neqsys.names, ab)), params))

varied_data = np.linspace(.5, 3)
xres, sols = neqsys.solve_series(x0, params, varied_data, 'p')

neqsys.plot_series(xres, varied_data, 'p')
plt.legend()
fig = plt.gcf()

"""
Explanation: The equations form a two dimensinoal non-linear problem (the problem is originally from here)
End of explanation
"""


def solve_custom(solver, plot_attr=None, **kwargs):
    ncols = 3 if plot_attr is None else 4
    ab, sol = neqsys.solve(x0, params, attached_solver=solver, ftol=1e-5, **kwargs)
    x_history = np.array(solver.history_x)
    plt.figure(figsize=(15, 3))
    plt.subplot(1, ncols, 1)
    plt.plot(x_history[:, 0], x_history[:, 1]); plt.xlabel('x0'), plt.ylabel('x1')
    plt.subplot(1, ncols, 2)
    plt.plot(neqsys.rms(x_history, [1])); plt.xlabel('iteration'), plt.ylabel('RMS(residuals)')
    plt.subplot(1, ncols, 3)
    plt.semilogy(range(15, len(x_history)), neqsys.rms(x_history[15:], [1]))
    plt.xlabel('iteration'); plt.ylabel('RMS(residuals)')
    if plot_attr is not None:
        plt.subplot(1, ncols, 4)
        plt.plot(np.asarray(getattr(solver, plot_attr)))
        plt.ylabel(plot_attr)
        plt.xlabel('iteration')
    plt.tight_layout()
    return sol

"""
Explanation: That was easy, let's use this system of equations to explore different algortihms.
Custom solvers
We can also use a custom solver, there are examples of simple (not production quality) solvers in pyneqsys.solvers. But they can be useful for exploring robustness of different classes of solvers for a given problem.
End of explanation
"""


from pyneqsys.solvers import DampedGradientDescentSolver
solve_custom(DampedGradientDescentSolver(1, 0))  # Undamped

"""
Explanation: Let's start with Gradient descent
End of explanation
"""


solve_custom(DampedGradientDescentSolver(.5, .5))  # (under-)damped

solve_custom(DampedGradientDescentSolver(.05, .05))  # over-damped

solve_custom(DampedGradientDescentSolver(.3, .1))   # (close to) optimally  damped

"""
Explanation: That didn't go too well, the name of the class may give a hint of a possible solution. Let's damp the steps:
End of explanation
"""


from pyneqsys.solvers import AutoDampedGradientDescentSolver
solve_custom(AutoDampedGradientDescentSolver(.1, .1, 8, .3), 'history_damping')

solve_custom(AutoDampedGradientDescentSolver(.1, .1, 6, .3), 'history_damping')

solve_custom(AutoDampedGradientDescentSolver(.1, .1, 4, .3), 'history_damping')

solve_custom(AutoDampedGradientDescentSolver(.1, .1, 2, .3), 'history_damping')

"""
Explanation: So we should strive to limit the osciallatory behaviour, let's see if we can achieve that algorithmically:
End of explanation
"""


from pyneqsys.solvers import PolakRibiereConjugateGradientSolver as PRCG
solve_custom(PRCG(5), 'history_sn', maxiter=50)

"""
Explanation: Another way to overcome the criss-cross behaviour of gradient descent is to use a generalized conjugate gradient solver, let's see how such a solver performs (note that there are a lot of function calls in the line searches).
End of explanation
"""


