from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import numpy as np
import math as mt

 
fig = plt.figure()
ax = Axes3D(fig, azim = -128, elev = 43)
s = .05
X = np.arange(-2, 2.+s, s) #arange(start,finish,increment), stores resulting vector in X
Y = np.arange(-2, 3.+s, s)
X, Y = np.meshgrid(X, Y)
# Rosenbrock Function
Z = (1.-X)**2 + 100.*(Y-X*X)**2
ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1, norm = LogNorm(), cmap = cm.jet)

# Rosenbrock funcion evaluation
def rosenbrock(ss):
	return (1.-ss[0])**2 + 100.*(ss[1]-ss[0]*ss[0])**2

def snext(s1,s2,s3): #convention s1 is largest
	smean = [0,0]
	stemp = [0,0]
	stemp1 = [0,0]
	snew = [0,0]

	smean[0] = (s2[0]+s3[0])/2 
	smean[1] = (s2[1]+s3[1])/2
	
	stemp[0] = 2*smean[0] - s1[0]
	stemp[1] = 2*smean[1] - s1[1]

	if rosenbrock(stemp) < rosenbrock(s1):
		stemp1[0] = 2*stemp[0] - s1[0]
		stemp1[1] = 2*stemp[1] - s1[1]
		if rosenbrock(stemp1) < rosenbrock(s1):
			snew = stemp1
		else:
			snew = stemp
	else:
		snew[0] = (s1[0] + smean[0])/2
		snew[1] = (s1[1] + smean[1])/2
	return snew
		

# Simplex starting point
s1 = [-1,-1]
s2 = [-1.1,-1]
s3 = [-1,-1.1]

path = np.zeros((2,1000))
pathval = np.zeros(1000)


converged = False
for i in range(1,1000): 
	if rosenbrock(s1) > rosenbrock(s2) and rosenbrock(s1) > rosenbrock(s3):
		s1 = snext(s1,s2,s3)
		path[:,i] = s1
		pathval[i] = rosenbrock(s1)
	elif rosenbrock(s2) > rosenbrock(s1) and rosenbrock(s2) > rosenbrock(s3):
		s2 = snext(s2,s3,s1)
		path[:,i] = s2
		pathval[i] = rosenbrock(s2)
	elif rosenbrock(s3) > rosenbrock(s1) and rosenbrock(s3) > rosenbrock(s2):
		s3 = snext(s3,s1,s2)
		path[:,i] = s3
		pathval[i] = rosenbrock(s3)
	if mt.fabs(pathval[i]- pathval[i-1]) < 1e-14:
		break

print "x = %f" %path[0,i]
print "y = %f" %path[1,i]
print "Fmin = %f" %pathval[i]
	

 
ax.plot(path[0,:],path[1,:],pathval[:],'bo-')
 

plt.xlabel("x")
plt.ylabel("y")
 
plt.savefig("Rosenbrock function.svg")
 
plt.show()