import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from ipywidgets import interact

"""
Explanation: <center>
    <img src="http://sct.inf.utfsm.cl/wp-content/uploads/2020/04/logo_di.png" style="width:60%">
    <h1> INF-285 - Computación Científica </h1>
    <h2> Newton's Method in $\mathbb{R}^n$ </h2>
    <h2> <a href="#acknowledgements"> [S]cientific [C]omputing [T]eam </a> </h2>
    <h2> Version: 1.05</h2>
</center>
<div id='toc' />

Table of Contents

Newton's method
Python Modules and Functions
Acknowledgements
End of explanation
"""


f1 = lambda x: x[0]**2+x[1]**2-1
f2 = lambda x: x[1]-x[0]**2
F  = lambda x: np.array([f1(x),f2(x)], dtype=float)
J  = lambda x: np.array([[2*x[0], 2*x[1]],[-2*x[0], 1.0]], dtype=float)

"""
Explanation: <div id='newton' />

Newton's method
Back to TOC
Using a vectorized implementation
End of explanation
"""


# Here we use 'np.linalg.solve', but we could use LU or PALU! Or iterative methods!
NewtonStep = lambda xi: xi-np.linalg.solve(J(xi),F(xi))

"""
Explanation: The Newton Step takes advantage of the vectorized implementation!
End of explanation
"""


n_delta = 50
x = np.linspace(-1.5, 1.5, n_delta)
# We could have used 'x' since it is the same, but for completeness we will define 'y'
y = np.linspace(-1.5, 1.5, n_delta)
X, Y = np.meshgrid(x, y)
Z1 = np.zeros_like(X)
Z2 = np.zeros_like(X)

for i,xi in enumerate(x):
    for j,yj in enumerate(y):
        Z1[j,i] = f1([xi,yj])
        Z2[j,i] = f2([xi,yj])

"""
Explanation: The next cell of code is just needed for plotting purposes.
End of explanation
"""


plt.figure()
CS1 = plt.contour(X, Y, Z1,levels=[0])
CS2 = plt.contour(X, Y, Z2,levels=[0])
plt.grid()
plt.axis('equal')
plt.title(r'Newton $\mathbb{R}^n$')
plt.show()

"""
Explanation: Here we plot the curves we will be intersecting.
End of explanation
"""


def Show_Newton(x0=1.2,y0=0.3,n=0):
    plt.figure()
    CS1 = plt.contour(X, Y, Z1,levels=[0])
    CS2 = plt.contour(X, Y, Z2,levels=[0])
    plt.grid()
    plt.axis('equal')
    plt.title(r'Newton $\mathbb{R}^n$')
    plt.plot(x0,y0,'rx')
    x_previous = np.array([x0,y0])
    print('Initial guess: [%.10f, %.10f]' % (x0,y0))
    for i in np.arange(n):
        x_next=NewtonStep(x_previous)
        x1,y1 = x_next
        plt.plot(x1,y1,'rx')
        plt.plot([x0, x1],[y0, y1],'r')
        x0=x1
        y0=y1
        x_previous = x_next
        print('Iteration %d  : [%.10f, %.10f]' % (i+1,x1,y1))
    plt.show()

interact(Show_Newton,x0=(-1.4,1.4,0.1),y0=(-1.4,1.4,0.1), n=(0,100,1))

"""
Explanation: Here we implement the Newton's method in higher dimension in a widget.
Most of the code is for plotting and showing the outcome, but the core, i.e. the Newton steps, is just 1 line of code!
End of explanation
"""


