import numpy as np
from numpy import *
from numpy import linalg as la
import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
# get_ipython().magic(u'matplotlib inline')

func_eval = 0

def f(x,y):
    global func_eval
    func_eval += 1
    f = (1-x)**2 + 100*(y-x**2)**2
    return f

def golden_line(x0,y0,dir=[1,0]):
    b0 = 0
    a0 = -3; a2 = 3
    a1 = 0.618034*(a2-a0)+a0
    a3 = 0.618034*(a1-a0)+a0
    x = np.dot(dir,[a3,b0])+x0
    y = np.dot(dir,[b0,a3])+y0
    last_f = f(x,y)
    last_x = x
    diff_x = 1
    while diff_x > 1e-6:
        x = np.dot(dir,[a3,b0])+x0
        y = np.dot(dir,[b0,a3])+y0
        x1 = np.dot(dir,[a1,b0])+x0
        y1 = np.dot(dir,[b0,a1])+y0
        # print str(x) + "\t" + str(y)
        if (f(x,y) < f(x1,y1)):
            a2 = a1
            a1 = a3
            a3 = 0.618034*(a1-a0)+a0
        else:
            a0 = a3
            a3 = a1
            a1 = a2-0.618034*(a2-a3)
        x = np.dot(dir,[a3,b0])+x0
        y = np.dot(dir,[b0,a3])+y0
        diff_x = abs(x - last_x)
        last_x = x
    da = pow((x-x0)**2 + (y-y0)**2,0.5)
    return x,y,da

pp = []
pp.append([-1, -1, f(-1,-1)])
xs = -1; ys = -1
d = zeros([2,3])
d[0,:] = [0,1,0]
d[1,:] = [0,0,1] 

x = xs; y = xs
x_last = xs; y_last = ys
a_min = -3; a_max = 3
diff_f = 1
it = 0
# for i in range(100):
while diff_f > 1e-5:
    print "dir: " + str(d[0,1:]) + "\t@ " + str([x,y])
    x,y,da = golden_line(x, y, d[0,1:])
    pp.append([x,y,f(x,y)])
    if it == 0:
        d[0,0] = da
    print "coord: " +str([x,y]) + "\tmoved: " + str(da)

    print "dir: " + str(d[1,1:]) + "\t@ " + str([x,y])
    x,y,da = golden_line(x, y, d[1,1:])
    pp.append([x,y,f(x,y)])
    d_new = [x-x_last, y-y_last]
    d_new = d_new/la.norm(d_new)
    diff_f = abs(f(x,y)-f(x_last,y_last))
    x_last = x; y_last = y
    if (da > d[0,0]):
        d[1,0] = da
        d[1,1:] = d_new
    else:
        d[1,:] = d[0,:]
        d[0,0] = da
        d[0,1:] = d_new
    print "coord: " +str([x,y]) + "\tmoved: " + str(da)
    it += 1

print "Done!"
print "Minimum @ " + str([x,y])
print "Iterations: " + str(it)
print "Function Evaluations: " + str(func_eval)

x = arange(-3,3,0.1)
y = arange(-3,3,0.1)
xx, yy = meshgrid(x, y)
z = f(xx, yy)

fig = plt.figure()
ax = Axes3D(fig)
ax.set_zlim(0,amax(z))
# ax = plt.axes(xlim=(0,NX), ylim=(0,NY))
# wireframe = ax.plot_wireframe(xx, yy, z)#,color="black",rstride=1, cstride=1)
contour = ax.contourf(xx,yy,z,200)
# ax.plot(pp0[:,0],pp0[:,1],zz0,'k')
# ax.plot([-1,x3,x3],[-1,-1,y3],[f(-1,-1),f(x3,-1),f(x3,y3)],'.r')

pp = reshape(pp,shape(pp))
ax.plot(pp[:,0],pp[:,1],pp[:,2],'.r',)
# ax.plot(pp1[:,0],pp1[:,1],zz1,'.k')
# ax.plot(pp2[:,0],pp2[:,1],zz2,'.k')
# contour = ax.contourf(xx,yy,z,200)
plt.show()