import matplotlib as mpl
import matplotlib.animation as animation
import matplotlib.collections as mcoll
import matplotlib.path as mpath
import matplotlib.pyplot as plt
import mpmath
import numpy as np

mpl.use("Agg")

# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=60, metadata=dict(artist='Me'), bitrate=4000)

minutes = 10
y = np.linspace(0, 90 * minutes, 3600 * minutes) # 1/40 ratio 

z_real = np.zeros_like(y)
z_imag = np.zeros_like(y)

z_real = []
z_imag = []
z_max = 0.0

ims = []
fig = plt.figure()

with writer.saving(fig, "output.mp4", dpi=300):
    for j, y_j in enumerate(y):
        print("%08d" % j, end='\r')

        z_ij = mpmath.zeta(complex(0.5, y_j))
    
        z_real += [z_ij.real]
        z_imag += [z_ij.imag]
        z_max = max(z_max, max(abs(z_ij.real), abs(z_ij.imag)))
    
        color_array = np.asarray(np.logspace(0.0, 1.0, len(z_real), base=j**2), dtype=float)
    
        points = np.array([z_real, z_imag]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        cmap = plt.get_cmap('gnuplot')
        norm = None
    
        plt.gca().add_collection(mcoll.LineCollection(segments, array=color_array, cmap=cmap, norm=norm, linewidth=3, alpha=1.0))
    
        ax_lim = np.float(1.1 * z_max)
    
        plt.plot([z_ij.real], [z_ij.imag], marker='*')
        plt.axis('off')
        plt.gca().set_xlim(-ax_lim, ax_lim)
        plt.gca().set_ylim(-ax_lim, ax_lim)

        writer.grab_frame()
        plt.cla()