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))

# 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 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.01
	while l(v0, v2) > min_length:
		
		# Evaluate function
		f0 = Rosenbrock(v0)
		f1 = Rosenbrock(v1)
		f2 = Rosenbrock(v2)
		f_evals += 3
		
		# 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 set, along system axes
D = array([[1, 0], [0, 1]]).astype(float)

# Minimize along directions, and replace most influential vector with total change
v0_ori = v0

min_step = 0.1
for i in range(80):
	for i in arange(2):
		v1 = line_minimization(v0_ori, D[0])
		plot_line(v0_ori, 0, v1, 0)
		v2 = line_minimization(v1, D[1])
		plot_line(v1, 0, v2, 0)
	
	if l(v0_ori, v2) < min_step:
		print i
		break
		
	# Overwrite most influential vector with total change
	D_new = (v2 - v0_ori)/l(v0_ori, v2)
	dots = dot(D, D_new)
	D[argsort(dots)[-1]] = D_new
	
	# Set new minimum point
	v0_ori = v2

# 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()