# -*- coding: utf-8 -*-
"""
Created on Mon Oct 18 16:03:25 2021

@author: Martin
"""

import numpy as np
print('numpy: '+np.version.full_version)
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D 
import matplotlib.animation as animation
import matplotlib
from scipy.special import sph_harm
from matplotlib import cm
print('matplotlib: '+matplotlib.__version__)



params = [[1, 0, 0, 0], 
          [1, 0, 1, 0]] #l1, l2, m1, m2



fps = 20 # frame per sec
frn = 50 # frame number of the animation
N = 32

phi = np.linspace(0, np.pi, N)
theta = np.linspace(0, 2*np.pi, N)
phi, theta = np.meshgrid(phi, theta)
# 
# The Cartesian coordinates of the unit sphere
x = np.sin(phi) * np.cos(theta)
y = np.sin(phi) * np.sin(theta)
z = np.cos(phi)

def g(phase, l1, l2, m1, m2):
    a, b = 1, phase
    # Calculate the spherical harmonic Y(l,m) and normalize to [0,1]
    fcolors = np.abs(a * sph_harm(m1, l1, theta, phi) + b * sph_harm(m2, l2, theta, phi))**2
    fmax, fmin = fcolors.max(), fcolors.min()
    fcolors = (fcolors - fmin)/(fmax - fmin)
    return fcolors

phases = np.exp(-1j * np.linspace(0, 2*np.pi, frn))

phases2 = np.exp(-1j * np.linspace(0, 2*np.pi, frn) + np.pi*1j)

fcolors1 = np.zeros((N, N, frn))
fcolors2 = np.zeros((N, N, frn))

for i in range(frn):
    fcolors1[:, :, i] = g(phases[i], *params[0])
    fcolors2[:, :, i] = g(phases2[i], *params[1])

def update_plot(frame_number, fcolors, plot):
    # fcolors = g(phase) 
    plot[0].remove()
    plot[0] = ax1.plot_surface(x, y, z,  rstride=1, cstride=1, facecolors=cm.jet(fcolors1[:, :, frame_number]))
    plot[1].remove()
    plot[1] = ax2.plot_surface(x, y, z,  rstride=1, cstride=1, facecolors=cm.jet(fcolors2[:, :, frame_number]))



fig = plt.figure(1)
ax1 = fig.add_subplot(121, projection='3d')
ax2 = fig.add_subplot(122, projection='3d')
for ax in (ax1, ax2):
    ax.view_init(elev=0, azim=0)
    pass

plot = [ax1.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=cm.seismic(fcolors1[:, :, 0])),
        ax2.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=cm.seismic(fcolors2[:, :, 0]))]
ax1.set_axis_off()
ax2.set_axis_off()

ax1.view_init(25, 0)
ax2.view_init(25, 0)

# ax1.set_title(fr'$ \left|   {params[0][0]},  {params[0][2]} \right> +' +'e^{i\omega_0 t}'+ fr'\left|   {params[0][1]},  {params[0][3]} \right>$ ')
# ax2.set_title(fr'$ \left|   {params[1][0]},  {params[1][2]} \right> -' +'e^{i\omega_0 t}'+ fr'\left|  {params[1][1]}, {params[1][3]} \right>$ ')
# fig.suptitle(r'$ \left|l_1, m_1  \right> +e^{i\omega_0 t}\left|l_2, m_2  \right>  $')


ani = animation.FuncAnimation(fig, update_plot, frn, fargs=(fcolors1, plot), interval=1000/fps)

# fn = complete file addres
# ani.save(fn+'.mp4',writer='ffmpeg',fps=fps)
# ani.save(fn+'.gif',writer='imagemagick',fps=fps)



#%%