import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib import cm

def Y(l,m,theta,phi):
    if l==0:
        return np.ones(theta.shape)*0.707
    elif l==1:
        if m==0:
            return np.cos(theta)
        elif m==1:
            return np.sin(theta)*np.cos(phi)
        elif m==-1:
            return np.sin(theta)*np.sin(phi)
    elif l==2:
        if m==0:
            return 1.5*(np.cos(theta)**2) - 0.5
        elif m==1:
            return np.sin(2*theta)*np.cos(phi)
        elif m==-1:
            return np.sin(2*theta)*np.sin(phi)
        elif m==2:
            return (np.sin(theta)**2)*np.cos(2*phi)
        elif m==-2:
            return (np.sin(theta)**2)*np.sin(2*phi)
    elif l==3:
        if m==0:
            c = np.cos(theta)
            return (2.5*c**2 - 1.5)*c
        elif m==1:
            return (5*(np.cos(theta)**2) - 1)*np.sin(theta)/1.377*np.cos(phi)
        elif m==-1:
            return (5*(np.cos(theta)**2) - 1)*np.sin(theta)/1.377*np.sin(phi)
        elif m==2:
            return np.sqrt(27)/2 * np.cos(theta)*(np.sin(theta))**2 * np.cos(2*phi)
        elif m==-2:
            return np.sqrt(27)/2 * np.cos(theta)*(np.sin(theta))**2 * np.sin(2*phi)
        elif m==3:
            return np.sin(theta)**3 * np.cos(3*phi)
        elif m==-3:
            return np.sin(theta)**3 * np.sin(3*phi)
    elif l==4:
        if m==0:
            c = np.cos(theta)**2
            return (35*(c**2) - 30*c + 3)/8
        elif m==1:
            c = np.cos(theta)
            return (7*c**2 - 3)/1.056*c*np.sin(theta)*np.cos(phi)
        elif m==1:
            c = np.cos(theta)
            return (7*c**2 - 3)/1.056*c*np.sin(theta)*np.cos(phi)
        elif m==2:
            return (7*(np.cos(theta)**2) - 1)*np.sin(theta)**2 /1.284 *np.cos(2*phi)
        elif m==-2:
            return (7*(np.cos(theta)**2) - 1)*np.sin(theta)**2 /1.284 *np.sin(2*phi)
        elif m==3:
            return 16/np.sqrt(27)*np.cos(theta) * np.sin(theta)**3 * np.cos(3*phi)
        elif m==-3:
            return 16/np.sqrt(27)*np.cos(theta) * np.sin(theta)**3 * np.sin(3*phi)
        elif m==4:
            return np.sin(theta)**4 * np.cos(4*phi)
        elif m==-4:
            return np.sin(theta)**4 * np.sin(4*phi)

def main():
	phi = np.linspace(0, 2* np.pi, 100)
	theta = np.linspace(0, np.pi, 100)
	phi, theta = np.meshgrid(phi, theta)

	fig = plt.figure(figsize=(12,12))
	plt.subplots_adjust(wspace=0, hspace=0)
	for l in range(5):
		for m in range(l+1):        
			Ylm = Y(l,m,theta,phi)
			r = np.abs(Ylm)

			x = r*np.sin(theta)*np.cos(phi)
			y = r*np.sin(theta)*np.sin(phi)
			z = r*np.cos(theta)
			
			ax = fig.add_subplot(5, 5, 5*l+m+1, projection='3d')
			ax.plot_surface(x, y, z, facecolors = cm.bwr(Ylm), antialiased=True)
			ax.set_xlim3d(-1,1)
			ax.set_ylim3d(-1,1)
			ax.set_zlim3d(-1,1)
			ax.view_init(30, -45)
			ax.axis('off')
	plt.savefig('spherical_harmonics.png')

if __name__=='__main__':
	main()