import numpy as np
import matplotlib.pyplot as plt
import imageio

plt.rcParams['text.usetex'] = True
plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'

a = 1.0
b = 3.0
f = lambda x: x**2 - 2.0
xs = np.arange(a, b, step=0.01)
ys = f(xs)

# Newton-Raphson method
N = 6
F = lambda x: x - (x**2 - 2.0)/(2.0*x)
x = [ 2.0 ]
for n in range(1, N + 1):
    x.append(F(x[n - 1]))
y = [ f(xn) for xn in x ]

T = 4*N
alpha = 0.3
blue = 'tab:blue'
orange = 'tab:orange'

for t in range(T - 3):
    n = t//4

    # graph y = x^2 - 2
    plt.plot(xs, ys, color=blue)
    plt.xlim(1.35, 2.05)
    plt.ylim(-0.2, 2.2)
    plt.xlabel(r'$x$')
    plt.ylabel(r'$y$')
    plt.title(r'Newton-Raphson method to find a zero of $f(x) = x^2 - 2$')

    # line y = 0
    plt.plot(xs, 0*ys, color=blue)

    # main animation
    for m in range(n + 1):
        # points x^(n)
        plt.plot(x[m], 0.0, 'o', color=orange, alpha=alpha)

    for m in range(n):
        # vertical line segments + points f(x^(n)) + tangent line segments
        plt.plot([ x[m], x[m] ], [ 0, y[m] ], color=orange, alpha=alpha)
        plt.plot(x[m], y[m], 'o', color=orange, alpha=alpha)
        plt.plot([ x[m + 1], x[m] ], [ 0, y[m] ], color=orange, alpha=alpha)

    if t % 4 == 0:
        # point x^(n) + text label
        plt.plot(x[n], 0.0, 'o', color=orange)
        plt.text(x[n] - 0.01, -0.15, r'$x^{(' + f'{n}' + r')}$')

    elif t % 4 == 1:
        # vertical line segment
        plt.plot([ x[n], x[n] ], [ 0, y[n] ], color=orange)

    elif t % 4 == 2:
        # point f(x^(n)) + text label (+ vertical line segment)
        plt.plot(x[n], y[n], 'o', color=orange)
        plt.text(x[n] - 0.04, y[n] + 0.1, r'$f(x^{(' + f'{n}' + r')})$')
        plt.plot([ x[n], x[n] ], [ 0, y[n] ], color=orange, alpha=alpha)

    elif t % 4 == 3:
        # tangent line segment (+ point f(x^(n)) + vertical line segment)
        plt.plot([ x[n + 1], x[n] ], [ 0, y[n] ], color=orange)
        plt.plot(x[n], y[n], 'o', color=orange, alpha=alpha)
        plt.plot([ x[n], x[n] ], [ 0, y[n] ], color=orange, alpha=alpha)

    # numerical results
    plt.text(1.407, 1.8, r'$\sqrt{2} = 1.414213562373095$')
    for m in range(n + 1):
        plt.text(1.400, 1.7 - 0.1*m, r'$x^{(' + f'{m}' + r')} = ' + f'{round(x[m], 15)}'.ljust(2 + 15, '0') + r'$')

    plt.savefig(f'{t}.png', dpi=300)
    plt.close()

filenames = [ f'{t}.png' for t in range(T - 3) ]
with imageio.get_writer('newton_raphson_method.gif', mode='I', fps=1) as writer:
    for filename in filenames:
        image = imageio.imread(filename)
        writer.append_data(image)