from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt
from numpy import *

# Define surface function
def Rosenbrock(v):
	x = v[0]
	y = v[1]
	return (1-x)**2 + 100*(y - x**2)**2

# Prepare figure
fig = plt.figure()
ax = fig.gca(projection='3d')

# Initialize number of function evaluations
f_evals = 0

# Define function to plot lines
def plot_line(v0, f0, v1, f1):
	A = vstack([v0, v1])
	X = A[:, 0]
	Y = A[:, 1]
	Z_surf = array([f0, f1])
	Z_proj = array([0, 0])
	verts_surf = [zip(X, Y, Z_surf)]
	ax.scatter(X, Y, Z_surf)
	ax.add_collection3d(Line3DCollection(verts_surf, colors='k', linewidths=0.5))
	verts_proj = [zip(X, Y, Z_proj)]
	ax.scatter(X, Y, Z_proj)
	ax.add_collection3d(Line3DCollection(verts_proj, colors='k', linewidths=0.5))

# Initialize line. Start near (x0, y0) = (-1, -1)
v0 = array([-1, -1]).astype(float)
f0 = Rosenbrock(v0)
f_evals += 1

# Computes the distance between points v1 and v2
def l(v1, v2):
	return sqrt(dot(v2 - v1, v2 - v1))

# Return local minimum of function along a line
# Line begins at v0 and extends parallel to dir
def line_minimization(v0, dir):
	
	# Declare number of function evaluations as global variable
	global f_evals
	
	# Define golden ratio
	phi = (1 + sqrt(5))/2
	shrink = phi/(phi + 1)
	grow = 1/shrink
	
	# Generate initial bracket
	v1 = v0 + dir*shrink
	v2 = v0 + dir
	
	# Bracket minimum
	min_length = 0.001
	while l(v0, v2) > min_length:
		
		# Evaluate function
		f0 = Rosenbrock(v0)
		f1 = Rosenbrock(v1)
		f2 = Rosenbrock(v2)
		f_evals += 3
		
		# Plot line
		plot_line(v0, f0, v1, f1)
		
		# Determine location of minimum
		o = argsort(array([f0, f1, f2]))
		
		# Extend or shrink line, use temporary point v3
		if o[0] == 0:
			# Minimum value at v0, extend by golden ratio to the left
			v3 = v0 - phi*(v1 - v0)
			# Overwrite previous values
			v2[:] = v1[:]
			v1[:] = v0[:]
			v0[:] = v3[:]
			continue
		elif o[0] == 1:
			# Minimum value at v1, build internal point
			if l(v1, v2) < l(v1, v0):		# Build point between v0 and v1
				v3 = v0 + shrink*(v1 - v0)
				f3 = Rosenbrock(v3)
				f_evals += 1
				# Overwrite previous values
				if f3 < f1:
					v2[:] = v1[:]
					v1[:] = v3[:]
					continue
				else:
					v0[:] = v3[:]
					continue
			else:							# Build point between v1 and v2
				v3 = v2 - shrink*(v2 - v1)
				f3 = Rosenbrock(v3)
				f_evals += 1
				# Overwrite previous values
				if f3 < f1:
					v0[:] = v1[:]
					v1[:] = v3[:]
					continue
				else:
					v2[:] = v3[:]
		else:
			# Minimum value at v2, extend by golden ratio to the right
			v3 = v2 + phi*(v2 - v1)
			# Overwrite previous values
			v0[:] = v1[:]
			v1[:] = v2[:]
			v2[:] = v3[:]
			continue
	
	# Quadratic interpolation to estimate location of minimum
	return v1

# Define initial direction
D = array([1, 0])

# Conjugate gradient method
v0_ori = v0

def Rosen_grad(v0):
	x = v0[0]
	y = v0[1]
	partial_x = -2*(1-x) - 400*x*(y - x**2)
	partial_y = 200*(y - x**2)
	grad = array([partial_x, partial_y]).astype(float)
	grad /= sqrt(dot(grad, grad))
	return grad

for i in range(80):
	# Find new minimum from v0_ori along direction D
	v1 = line_minimization(v0_ori, D)
	
	# Calculate gradients at new and original points
	g0 = Rosen_grad(v0_ori)
	g1 = Rosen_grad(v1)
	f_evals += 2
	# Compute gamma factor, to determine new direction
	gamma = dot(g1, g1)/dot(g0, g0)
	# Calculate new (conjugate) direction
	D1 = g1 + gamma*D
	
	# Set new minimum point
	v0_ori = v1
	# Normalize new direction
	D = D1/dot(D1, D1)

# Plot Rosenbrock surface
X_Rosen = arange(-1.3, 1.3, 0.05)
Y_Rosen = arange(-1, 1.2, 0.05)
X_Rosen, Y_Rosen = meshgrid(X_Rosen, Y_Rosen)
Z_Rosen = (1 - X_Rosen)**2 + 100*(Y_Rosen - X_Rosen**2)**2
surf_Rosen = ax.plot_surface(X_Rosen, Y_Rosen, Z_Rosen, rstride=1, cstride=1,
	cmap=cm.coolwarm, linewidth=0, antialiased=False, alpha = 0.3)

# Adjust axes
ax.set_zlim(0, 600)
ax.zaxis.set_major_locator(LinearLocator(5))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.0f'))

# Report minimum
print 'Minimum location', v0_ori, '\nMinimum value', Rosenbrock(v0_ori), '\nNumber of function evaluations', f_evals

# Render plot
plt.show()