import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import time

## Prepare initial data
a = 20
d = 20
tm0 = 10
t_m = tm0 * np.pi/180
def phi(t):
    return np.pi/4 + (6 + 0.1*t) * t
def r(t):
    return d/2 * (1 - phi(t)/phi(t_m))
def Fp(x,y,t):
    return -1/np.sqrt((x + r(t)*np.cos(phi(t)))**2 + (y + r(t)*np.sin(phi(t)))**2)                                                            \
           -1/np.sqrt((x - r(t)*np.cos(phi(t)))**2 + (y - r(t)*np.sin(phi(t)))**2)     
X = np.arange(-1*a, a)
Y = np.arange(-1*a, a)
x,y=np.meshgrid(X,Y)

## generate image data
imageList = []
for t in range(tm0):
    imageList.append(Fp(x,y,t*np.pi/180))

## Create animation and video from 2D images
fig = plt.figure(figsize=(10,10))
ims = []
for i in range(len(imageList)):
    im = plt.pcolormesh(imageList[i], animated = True)
    ims.append([im])
    
ani1 = animation.ArtistAnimation(fig, ims, blit=True, repeat_delay=2000)
plt.colorbar()
t1=time.time()
ani1.to_html5_video()
t2=time.time()
print("2D image to video took: ", t2 - t1)

## Create animation and video from 3D images
t1 = time.time()
fig = plt.figure()
ax = Axes3D(fig)
ims = []
for i in range(len(imageList)):
    im = ax.plot_surface(x,y,imageList[i], antialiased=False, animated=True)
    ims.append([im])
    
ani = animation.ArtistAnimation(fig, ims, blit=True, repeat_delay=2000)
## plt.show()
t2 = time.time()
print("3D animation creation took: ", t2 - t1)
ani.to_html5_video()
t3=time.time()
print("3D animation to video took:", t3 - t2)