"""Generate the graph of the function $f(x,y) = -(\cos^2x + \cos^2y)^2$."""

import numpy as np
from pylab import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

# Thanks to CT Zhu. See the link as follows
# http://stackoverflow.com/questions/22867620/putting-arrowheads-on-vectors-in-matplotlibs-3d-plot
class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        FancyArrowPatch.draw(self, renderer)

# draw the frames
ax = Axes3D(figure(), azim = -135, elev = 45)
X = arange(-pi/2, pi/2, pi/100)	# set ranges
Y = arange(-pi/2, pi/2, pi/100)
X, Y = meshgrid(X, Y)
Z = -(cos(X)**2+cos(Y)**2)**2	# give the function
ax.plot_wireframe(X, Y, Z, rstride = 10, cstride = 10)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel('f(x,y)')
ax.set_xlim([-pi/2,pi/2])
ax.set_ylim([-pi/2,pi/2])
ax.set_zlim([-4,0])

# draw the contoured quiver
X1 = arange(-pi/2, pi/2, pi/20)	# start of the arrows
Y1 = arange(-pi/2, pi/2, pi/20)
for x in X1:
	for y in Y1:
		ax.add_artist(Arrow3D([x,x+4*sin(x)*cos(x)*(cos(x)**2+cos(y)**2)*0.06],\
		[y,y+4*sin(y)*cos(y)*(cos(x)**2+cos(y)**2)*0.06],\
		[-3.9999,-4.0000],lw=1, mutation_scale=5,arrowstyle="->", color="r"))

plt.savefig("gradient_vis.svg")