# -*- coding: utf-8 -*-
"""
date of creation: Fri Mar 25 16:20:46 2022
-------------------------------------------------
@author : Manuel Cattaneo
@contact: cattaneo.manuel@hotmail.com
-------------------------------------------------    

Description:
    Program used to play around with the spherical harmonics.

Modules:
"""
import numpy as np
# plot stuff
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d  # to plot 3d
import matplotlib.gridspec as gs              # for grid layout
# spherical harmonics
from scipy.special import sph_harm
"""
------------------------------------------------------------------------------
"""

if __name__ == '__main__':
    # grid spec
    row       = 3
    col       = 2*row-1                     
    grid_spec = gs.GridSpec(row,col)
    fig       = plt.figure()
    
    # creating spherical domain
    N = 50
    theta, phi = np.meshgrid(np.linspace(0, 2*np.pi, N), # theta in [0, 2pi]
                             np.linspace(0,   np.pi, N)) # phi   in [0,  pi]
    
    import matplotlib
    # m < n
    for y_grid in range(0,row):
        for i,x_grid in enumerate(range(col//2-y_grid, col-col//2+y_grid)):
            # real(...) to obtain only the cos(.) part. This is done in order   
            # to have the same functions as the ones at p.518 chapter 12.2.
            m = y_grid
            n = m+i+1
            r = np.real( sph_harm(m, n, theta, phi) ) # Y_n^m 
            # coordinates transformation
            x = r*np.cos(theta)*np.sin(phi)
            y = r*np.sin(phi)  *np.sin(theta)
            z = r              *np.cos(phi)
            
            # adding subplot
            ax = fig.add_subplot(grid_spec[y_grid,x_grid], projection='3d')
            ax.set_axis_off()
            ax.set_title(rf'$Y_{n}^{m}(\theta, \phi)$')
            
            # colors
            color_dim = r
            color_min, color_max = color_dim.min(), color_dim.max()
            norm = matplotlib.colors.Normalize(color_min,color_max)
            m = plt.cm.ScalarMappable(norm, cmap='jet')
            m.set_array([])
            fc = m.to_rgba(color_dim)
            
            # plot
            plot = ax.plot_surface(x, y, z, facecolors = fc,
                                   rstride=1, cstride=1, 
                                    # cmap=plt.get_cmap('jet'),
                                   linewidth=0, antialiased=False, alpha=0.5)
            
            
    plt.show()