"""
Animation of the Rosenbrock function of three variables
"""
import numpy as np
from scipy.optimize import rosen
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation

Axes3D
plt.rcParams["font.size"] = 10
plt.rcParams["mathtext.fontset"] = "cm"

x1, x2 = np.mgrid[-2.5:2.5:51j, -2.5:2.5:51j]
x3  = np.linspace(-2.5, 2.5, 21)
x3  = np.append(x3, x3[-2::-1])
fun = rosen([x1, x2, x3[0]])

fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111, projection='3d')
ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) 
ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) 
ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))

def data_gen(num):
    fun = rosen([x1, x2, x3[num]])
    ax.cla()
    surf = ax.plot_wireframe(x1, x2, fun/1000, linewidth=1, color="#3c3c3c",
                         rstride=51, cstride=51)
    contour = ax.contour(x1, x2, fun/1000, 20)
    ax.set_xlabel(r"$x_1$", fontsize=12)
    ax.set_ylabel(r"$x_2$", fontsize=12)
    ax.set_zlabel(r"$f(x_1, x_2, x_3)/10^3$", fontsize=12)
    ax.set_title(r"$x_3 = {}$".format(x3[num]), fontsize=12, y=1.08)
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-2.5, 2.5)
    ax.set_zlim(0, 16)
    ax.view_init(elev=45, azim=-45)
    return surf, contour

ani = animation.FuncAnimation(fig, data_gen, range(41))
ani.save("rosenbrock3.gif", writer='imagemagick')
plt.show()