import numpy as np
import mayavi.mlab as mlab
import  moviepy.editor as mpy
import os, time
from spherical_harmonic import spherical_harmonic
from latex_util import get_latex_img_data

width = 854
height = 526 #480
video_file = "spherical_harmonic_scene_1.mp4"
if os.path.exists(video_file):
    raise Exception("File already exists")

fig = mlab.figure(size=(width,height), bgcolor=(1.0,1.0,1.0), fgcolor=(0.,0.,0.))

# m0=4 ; m1=3 ; m2=2 ; m3=3 ; m4=6 ; m5=2 ; m6=6 ; m7=4
# m0=4 ; m1=2 ; m2=2 ; m3=2 ; m4=6 ; m5=4 ; m6=6 ; m7=2
m0=4 ; m1=8 ; m2=2 ; m3=6 ; m4=6 ; m5=4 ; m6=6 ; m7=2
x, y, z = spherical_harmonic(n=360.0, m0=m0, m1=m1, m2=m2, m3=m3, m4=m4, m5=m5, m6=m6, m7=m7)
mlab.mesh(x, y, z)
fig.scene.camera.roll(30)
#fig.scene.camera.yaw(5)
fig.scene.camera.elevation(90)
fig.scene.camera.zoom(1.4)

fm1 = r"$r = \sin^{m_1}(m_0\phi) + \cos^{m_3}(m_2\phi) + \sin^{m_5}(m_4\theta) + \cos^{m_7}(m_6\theta)$"
fm2 = r"$x = r \sin(\phi) \cos(\theta)$"
fm3 = r"$y = r \cos(\phi)$"
fm4 = r"$z = r \sin(\phi) \sin(\theta)$"
fm5 = r"$m_0=4, m_1=3, m_2=2, m_3=3, m_4=6, m_5=2, m_6=6, m_7=4$"
mlab.text(0.02, 0.90, fm1, figure=fig, width=0.45)
mlab.text(0.02, 0.83, fm2, figure=fig, width=0.15)
mlab.text(0.02, 0.76, fm3, figure=fig, width=0.1)
mlab.text(0.02, 0.69, fm4, figure=fig, width=0.15)
#mlab.text(0.02, 0.62, fm5, figure=fig, width=0.6)
fig.scene.render()

duration = 300.0 # duration of the animation in seconds
fps = 20 # frames per second
frames = (duration * fps)
frame = 0
zm = np.power(2.0, 1.0/frames)

def make_frame(t):
    global frame, m0, m1, m2, m3, m4, m5, m6, m7
    if t<2:
        return mlab.screenshot(antialiased=True)

    mlab.clf(figure=fig)
    frame += 1

    # scene 1
    #m0 += 0.0015
    #m2 += 0.0015
    #m4 -= 0.0015

    # scene 2
    #m0 += 0.001
    #m2 += 0.001
    #m4 += 0.001
    #m6 -= 0.001

    # scene 3
    #m2 += 0.0020
    #m4 += 0.0015
    #m6 -= 0.0015

    # scene 4
    #m2 += 0.0025
    #m4 += 0.0015
    #m6 -= 0.001

    # scene 5
    m0 += 0.0045
    m2 += 0.0035
    m4 += 0.0025
    m6 -= 0.001

    x, y, z = spherical_harmonic(n=360.0, m0=m0, m1=m1, m2=m2, m3=m3, m4=m4, m5=m5, m6=m6, m7=m7)
    mlab.mesh(x, y, z)
    fig.scene.camera.zoom(1.4)
    r = 2.5 * 360.0 / frames
    a = 5.0 * 360.0 / frames
    e = 10.0 / frames
    if t < 0.6 * duration:
        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.distance('auto')
    #fig.scene.camera.elevation(e)
    #fig.scene.camera.yaw(0.01)

    fmt = "frame=%d, time=%.2f, m0=%.3f, m1=%.3f, m2=%.3f, m3=%.3f, m4=%.3f, m5=%.3f, m6=%.3f, m7=%.3f"
    txt = fmt % (frame, t, m0, m1, m2, m3, m4, m5, m6, m7)
    t = mlab.text(0.02, 0.01, txt, figure=fig, width=0.95) #, color=(0.5, 0.5, 0.5), opacity=0.2)
    fig.scene.render()
    return mlab.screenshot(antialiased=True)

animation = mpy.VideoClip(make_frame, duration=duration)
#animation.resize(height=height, width=width)
#animation.write_gif("spherical_harmonic_2.gif", fps=fps)
animation.write_videofile(video_file, fps=fps)