# newton.py

# https://caam37830.github.io/book/04_functions/roots.html

import math
from frange import *

import matplotlib.pyplot as plt

EPS = 1E-09


def newtonRoot(fn, fnDer, x):
  """
  numerically approximate a root of fn using Newton's method
      fn - function
      fnDer - derivative of function fn
      x - starting point (default 1)
  returns:
      x - the root
      iters = a list of all x visited
      fs - fn(x) for each iteration
  """
  fx = fn(x)
  iters = []
  fs = []
  while abs(fx) > EPS:
    x = x - fx / fnDer(x)
    fx = fn(x)
    iters.append(x)
    fs.append(fx)
      
  return x, iters, fs


# ------------------------------------

if __name__ == "__main__":
  
  fn = lambda x : x**2 - 2
  fnDer = lambda x : 2*x
  
  # roots are +- sqrt(2)

  x , ix, iy = newtonRoot(fn, fnDer, 1)
  print("Root found starting at  1:", x)
  
  x, ix, iy = newtonRoot(fn, fnDer, -1)
  print("Root found starting at -1:", x)
  
  '''
  xmin, xmax = -2, 2
  xs = linspace(xmin, xmax, 100)
  x0 = 1
  plt.plot(xs, fn(xs), 'k-', label='fn') # plot the function
  plt.hlines(0, xmin, xmax, colors='gray', linestyles='dashed') # dashed line for y = 0
  plt.scatter(x0, fn(x0), c='k')
  plt.text(x0 + 0.05, fn(x0) - 0.1, r'$x_0$') # x0 label
  x1 = x0 - fn(x0)/fnDer(x0)
  plt.plot(xs, fnDer(x0)*(xs - x1), 'k--', label='tangent')
  plt.scatter(x1, 0, c='k')
  plt.text(x1 + 0.05, -0.2, r'$x_1$') # x1 label
  plt.xlim(xmin, xmax)
  plt.ylim(xmin, xmax)
  plt.legend()
  plt.show()
  '''