import numpy as np
import matplotlib.pyplot as plt


"""
Minimize Rosenbrock function
f(x0,x1) = 100 (x1-x0**2)**2 + (1-x0)**2
using steepest descent or Newtons method with Wolfe linesearch
"""

f  = lambda x: 100.0*(x[1]-x[0]**2)**2 + (1.0-x[0])**2
df = lambda x: np.array([-400.0*(x[1]-x[0]**2)*x[0] - 2.0*(1.0-x[0]),\
                         200.0*(x[1]-x[0]**2)])
d2f= lambda x: np.array([[-400.0*(x[1]-x[0]**2)+800.0*x[0]**2+2.0,-400.0*x[0]],\
                         [-400.0*x[0],200.0]])

def bisection_linesearch(f,df,x,p,alpha0,c1,c2):
    alpha_min, alpha_max = 0.0, np.inf
    alpha = alpha0
    fx   = f(x)
    dfx  = df(x)
    dfxp = dfx.dot(p)
    while alpha < 1.0E+10:
        #print(alpha,alpha_min,alpha_max)
        if f(x+alpha*p) > fx + alpha*c1*dfxp:
            # no sufficient decrease - too long step
            alpha_max = alpha
            alpha = 0.5*(alpha_max + alpha_min)
        elif df(x+alpha*p).dot(p) < c2*dfxp:
            # no curvature condition: too short step
            alpha_min = alpha
            if alpha_max == np.inf:
                alpha *= 2.0
            else:
                alpha = 0.5*(alpha_max + alpha_min)
        else:
            # we are done!
            return alpha
    raise ValueError('Steplength is way too long!')

# Wolfe condition parameters
c1 = 0.25
c2 = 0.5
# Stopping tolerance (stop when norm of the gradient is smaller than this)
tol = 1.0E-10
# Tolerance for switching to steepest descent
eps = 1.0E-08
# Starting point
#x0 = np.array([-1.2, 1])
x0 = np.random.randn(2)
# max no of iterations
Nmax = 10000
# Whether to use Newton's method or steepest descent
Newton = True
# start the loop
n = 0
alpha = 1.0

x = x0
while n<Nmax:
    grad = df(x)
    gradn= np.linalg.norm(grad,2)
    if gradn < tol:
        print("Success!")
        break
    if Newton:
        # Newton's method
        hess = d2f(x)
        p = -np.linalg.solve(hess,grad)
        if -grad.dot(p) <= eps*gradn*np.linalg.norm(p,2):
            # angle too large; use steepest descent direction instead
            p = -grad
            used_newton_step = False
        else:
            used_newton_step = True
    else:
        # steepest descent
        p = -grad
        used_newton_step = False
    alpha = bisection_linesearch(f,df,x,p,1.0,c1,c2)
    # print stuff
    print("Iter: %5d, f=%15.6e, |grad f|=%15.6e, alpha=%15.6e, Newton=%d" % \
          (n,f(x),gradn,alpha,used_newton_step))
    x += alpha*p
    n += 1

print("Iter: %5d, f=%15.6e, |grad f|=%15.6e, alpha=%15.6e" % \
      (n,f(x),gradn,alpha))