import matplotlib.pyplot as plt
import math

def newton(f_div_df, x):
    lo, hi = -float('inf'), +float('inf')
    apx = [x]
    while lo < x < hi:
        print(x)
        delta = f_div_df(x)
        if delta < 0.0:
           lo = max(lo, x)
        else:
           hi = min(hi, x)
        x -= delta
        apx.append(x)
    return apx

def visualize_newton(title, f, f_div_df, x0):
    apx = newton(f_div_df, x0)
    
    solution = apx[-1]
    iterations = len(apx) - 1
    samples = 100
    
    xmin, xmax = min(apx), max(apx)
    padding = (xmax - xmin) * 0.1
    xmin -= padding
    xmax += padding
    dx = (xmax - xmin) / (samples - 1)
    xs = [xmin + i * dx for i in range(samples)]
    ys = [f(x) for x in xs]

    plt.plot([xmin, xmax], [0, 0], 'k-', linewidth=0.5)
    plt.plot(xs, ys, label='$f(x)$')
    plt.plot([x for x in apx for _ in [1, 2]],
             [y for x in apx for y in [0, f(x)]],
             'r.-', label=f"Newton-Raphson, {iterations} iterations")
    plt.plot([x0], [0], 'go', label=f'Start $x_0 = {x0}$')
    plt.plot([solution], [f(solution)], 'ro', label=f'root $x = {solution}$')
    plt.legend(title=title)

# fraction

for fig, n in enumerate([.25, 49], start=1):
    plt.subplot(2, 2, fig)
    visualize_newton(
        r'$1 / %s$ : $f(x) = %s - 1/x$' % (n, n),
        lambda x: n - 1 / x,           # f(x)
        lambda x: (n * x - 1) * x,     # f(x) / f'(x)
        10.0**(-math.log(n, 10) // 1)  # closest smaller 10^k
    )

# square root

for fig, n in enumerate([.3, 5], start=3):
    plt.subplot(2, 2, fig)
    visualize_newton(
        r'$\sqrt{%s}$ : $f(x) = x^2 - %s$' % (n, n),
        lambda x: x * x - n,        # f(x)
        lambda x: (x - n / x) / 2,  # f(x) / f'(x)
        1.0 if n <= 1.0 else n
    )

plt.tight_layout()
plt.show()
