#! /usr/bin/env python3

# Scatter-plot of rotatorstates

import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.animation as animation

def calculate_data():
	Phi=lambda phi: 1/np.sqrt(2*np.pi)	# Azimutal part of wavefunction
	Th_s=lambda theta: 1/np.sqrt(2)		# Polar part of wavefunction            s-distribution
	Th_pz=lambda costheta: costheta*np.sqrt(1.5)	# Polar part of wavefunction    pz-distribution
	
	nphi=99
	vphi=np.linspace(0,2*np.pi,nphi+1)	# Vector of phi-values
	vPhi=np.zeros(nphi+1)			
	for i in range(nphi+1):
		vPhi[i]=Phi(vphi[i])		# Vector of Phi-values

	nth=99
	vcosth=np.linspace(-1,1,nth+1)		# Vector of cos(theta)-values
	vTh_s=np.zeros(nth+1)
	vTh_pz=np.zeros(nth+1)
	for i in range(nth+1):
		vTh_s[i]=Th_s(vcosth[i])		# Vector of Theta-values
		vTh_pz[i]=Th_pz(vcosth[i])		# Vector of Theta-values

	phi_interval=2*np.pi/nphi		# phi-interval
	costh_interval=2/nth			# Interval in cos(theta)
	# Cummulative sum of azimutal probability:
	CPphi=(np.cumsum(vPhi**2)-vPhi[0]**2)*phi_interval
	CPphi[-1]=1				# Make sure that it ends with 1
	CPth_s=(np.cumsum(vTh_s**2)-vTh_s[0]**2)*costh_interval
	CPth_s[-1]=1
	CPth_pz=(np.cumsum(vTh_pz**2)-vTh_pz[0]**2)*costh_interval
	CPth_pz[-1]=1


	NR_s=500					# Number of points
	CPphirand=np.random.rand(1,NR_s)		# Generate NR_s random valuesof integrated azimutal probability
	# Find the corresponding phi-values by interpolation:
	phirand_s=np.interp(CPphirand,CPphi,vphi)
	CPthrand=np.random.rand(1,NR_s)		# Generate NR_s random values og integrated polar probability
	# Find the corresponding values of cos(theta) by interpolation:
	costhrand_s=np.interp(CPthrand,CPth_s,vcosth)

	# Cartesian coordinates of the NR_s positions:
	r=4.0
	z_s=r*costhrand_s
	sinthrand_s=np.sqrt(1-costhrand_s**2)
	x_s=r*sinthrand_s*np.cos(phirand_s)
	y_s=r*sinthrand_s*np.sin(phirand_s)

	# Combine points in a 3xNR_s array
	data_s=np.empty((3,NR_s))
	data_s[0,:]=x_s
	data_s[1,:]=y_s
	data_s[2,:]=z_s

	NR_pz=1000					# Number of points
	CPphirand=np.random.rand(1,NR_pz)		# Generate NR_pz random valuesof integrated azimutal probability
	# Find the corresponding phi-values by interpolation:
	phirand_pz=np.interp(CPphirand,CPphi,vphi)
	CPthrand=np.random.rand(1,NR_pz)		# Generate NR_pz random values og integrated polar probability
	# Find the corresponding values of cos(theta) by interpolation:
	costhrand_pz=np.interp(CPthrand,CPth_pz,vcosth)

	# Cartesian coordinates of the NR_pz positions:
	r=4.0
	z_pz=r*costhrand_pz
	sinthrand_pz=np.sqrt(1-costhrand_pz**2)
	x_pz=r*sinthrand_pz*np.cos(phirand_pz)
	y_pz=r*sinthrand_pz*np.sin(phirand_pz)

	# Combine points in a 3xNR_pz array
	data_pz=np.empty((3,NR_pz))
	data_pz[0,:]=x_pz
	data_pz[1,:]=y_pz
	data_pz[2,:]=z_pz
	return data_s, data_pz

# Function to make two scatter-plot annimations of 3D data, with given title and markersize.
def animate_3d_scatter_rotating(data1,data2,title1,title2,msize):
	def update_lines(num, data, line, ax) :
		line.set_data(data[0:2, :num])
		line.set_3d_properties(data[2,:num])
		return line
	fig = plt.figure()
	ax1 = fig.add_subplot(1,2,1,projection='3d')
	ax2 = fig.add_subplot(1,2,2,projection='3d')
	line1 = ax1.plot(data1[0, 0:1], data1[1, 0:1], data1[2, 0:1],'o',markersize=msize)[0]
	line2 = ax2.plot(data2[0, 0:1], data2[1, 0:1], data2[2, 0:1],'o',markersize=msize)[0]
	ax1.set_title(title1)
	ax2.set_title(title2)
	ax1.view_init(elev=16, azim=30)
	ax2.view_init(elev=16, azim=30)

	cmax1=np.max(abs(data1))
	ax1.axis('off')
	ax1.plot3D([-cmax1,cmax1],[0,0],[0,0],'k')
	ax1.plot3D([0,0],[-cmax1,cmax1],[0,0],'k')
	ax1.plot3D([0,0],[0,0],[-cmax1,cmax1],'k')
	ax1.text(cmax1*1.05,0,0,'x')
	ax1.text(0,cmax1*1.05,0,'y')
	ax1.text(0,0,cmax1*1.05,'z')

	cmax2=np.max(abs(data2))
	ax2.axis('off')
	ax2.plot3D([-cmax2,cmax2],[0,0],[0,0],'k')
	ax2.plot3D([0,0],[-cmax2,cmax2],[0,0],'k')
	ax2.plot3D([0,0],[0,0],[-cmax2,cmax2],'k')
	ax2.text(cmax2*1.05,0,0,'x')
	ax2.text(0,cmax2*1.05,0,'y')
	ax2.text(0,0,cmax2*1.05,'z')


	Nframes1=len(data1[0,:])
	Nframes2=len(data2[0,:])
	line_ani1 = animation.FuncAnimation(fig, update_lines, frames=Nframes1, fargs=(data1, line1, ax1),interval=50, blit=False,repeat=False)
	line_ani2 = animation.FuncAnimation(fig, update_lines, frames=Nframes2, fargs=(data2, line2, ax2),interval=50, blit=False,repeat=False)
	plt.show()

if __name__=='__main__':
	title1='s distribution'
	title2='pz distribution'
	data1,data2=calculate_data()
	animate_3d_scatter_rotating(data1,data2,title1,title2,1)