# coding: utf-8

# In[1]:


import numpy as np
import scipy as sc
from moviepy.editor import VideoClip
from vispy import app, scene
from vispy.gloo.util import _screenshot
from IPython.display import clear_output, Image, display


# In[2]:


n = 7
alpha = 6
zero = sc.special.jn_zeros(n,alpha)[alpha-1]
duration = 64*np.pi/zero
print(zero,duration)


# In[3]:


X, Y = np.linspace(-1.024,1.024,256), np.linspace(-1.024,1.024,256)


# In[4]:


XX, YY = np.meshgrid(X,Y)


# In[5]:


RR = (XX**2+YY**2)**(0.5)
PP = np.arctan2(XX,YY)
ZZ = lambda t: sc.special.jn(n,zero*RR)*np.cos(n*PP)*np.cos(zero*t/32.0)*np.heaviside(1-RR,0)


# In[6]:


canvas = scene.SceneCanvas(keys='interactive', bgcolor='w')
view = canvas.central_widget.add_view()
view.camera = scene.TurntableCamera(up='z', fov=60)


# In[7]:


surface = scene.visuals.SurfacePlot(x=X, y=Y, z= ZZ(0),shading='smooth', color=(0.5, 0.5, 1, 1))
view.add(surface)
canvas.show()


# In[8]:


def make_frame(t):
    surface.set_data(z = ZZ(t))
    canvas.on_draw(None)
    return _screenshot((0,0,canvas.size[0],canvas.size[1]))[:,:,:3]


# In[10]:


animation = VideoClip(make_frame, duration=duration).resize(width=1024)
animation.write_gif("bessel_mode.gif", fps=20)


# In[11]:


clear_output(wait = True)
display(Image(url='bessel_mode.gif'))