import numpy as np
import mayavi.mlab as mlab
import  moviepy.editor as movedit
from spherical_harmonic import spherical_harmonic
from latex_util import get_latex_img_data

fm1 = r"$r = \sin^3(4\phi) + \cos^3(2\phi) + \sin^2(6\theta) + \cos^4(6\theta)$"
fm2 = r"$x = r \sin(\phi) \cos(\theta)$"
fm3 = r"$y = r \cos(\phi)$"
fm4 = r"$z = r \sin(\phi) \sin(\theta)$"

# MAKE A FIGURE WITH MAYAVI

fig = mlab.figure(size=(800,450), bgcolor=(1.0,1.0,1.0), fgcolor=(0.,0.,0.))

x, y, z = spherical_harmonic(n=360.0, m4=6)
mlab.mesh(x, y, z)

duration = 30.0 # duration of the animation in seconds (it will loop)
fps = 20 # frames per second
frames = (duration * fps)
zm = np.power(10.0, 1.0/frames)
fig.scene.camera.zoom(1.2)

# LaTeX formula of the parametric equations
mlab.text(0.02, 0.90, fm1, figure=fig, width=0.45) #, color=(0.5, 0.5, 0.5), opacity=0.2)
mlab.text(0.02, 0.82, fm2, figure=fig, width=0.15)
mlab.text(0.02, 0.74, fm3, figure=fig, width=0.1)
mlab.text(0.02, 0.66, fm4, figure=fig, width=0.15)

def make_frame(t):
    r = (duration/10.0) * 180.0 / frames
    a = (duration/10.0) * 360.0 / frames
    e = 10.0 / frames
    y = t * 20.0 / frames
    if t < duration/4 or 2*duration/4 < t < 3*duration/4:
        fig.scene.camera.zoom(zm)
    else:
        fig.scene.camera.zoom(1.0/zm)
    fig.scene.camera.roll(r)
    fig.scene.camera.azimuth(a)
    fig.scene.camera.elevation(e)

    fig.scene.render()
    return mlab.screenshot(antialiased=True)

animation = movedit.VideoClip(make_frame, duration=duration)
animation.write_gif("spherical_harmonic.gif", fps=fps)