#!/usr/bin/env python

import sys, os
import numpy as np
from scipy import stats
#from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import *
import matplotlib.pyplot as plt
import matplotlib as mpl

def CalculateRotationalMatrix( phi, theta, psi):

        RotationalMatrix = np.zeros((3,3))

        RotationalMatrix[0][0] =   -1.*np.sin(psi)*np.sin(phi)  +  np.cos(theta)*np.cos(phi)*np.cos(psi)
        RotationalMatrix[0][1] =       np.cos(phi)*np.sin(psi)  +  np.cos(theta)*np.cos(psi)*np.sin(phi)
        RotationalMatrix[0][2] =   -1.*np.cos(psi)*np.sin(theta)
        RotationalMatrix[1][0] =   -1.*np.sin(phi)*np.cos(psi)  -  np.cos(theta)*np.sin(psi)*np.cos(phi)
        RotationalMatrix[1][1]  =      np.cos(psi)*np.cos(phi)  -  np.cos(theta)*np.sin(phi)*np.sin(psi)
        RotationalMatrix[1][2] =       np.sin(psi)*np.sin(theta)
        RotationalMatrix[2][0] =       np.sin(theta)*np.cos(phi)
        RotationalMatrix[2][1] =       np.sin(theta)*np.sin(phi)
        RotationalMatrix[2][2] =       np.cos(theta)

        return RotationalMatrix


# J,K,M value
J = 1.
K = 0.
M = 1.

Npoints = 100
Mpoints = 100

mpl.rc("font", size=15 )

CH = np.array( [ 0., 0., 1. ] )

#  Setup first rotation: orient the molecule with J along z-axis, and projection
#  of J on molecular symmetry axis equal to K

alpha  = 0.  # Random rotation of molecule around figure axis (not needed)
beta   = np.arccos(  K / np.sqrt( J *( J + 1 ) ) )
gammas  = np.linspace( 0.,  2.* np.pi, Npoints )  # Random projection of J on (X,Y)-plane

#  Setup second rotation: orient the molecule in such a way that the projection
#  of J on z-axis is equal to M

phi   = 0. 
theta = np.arccos( M  / np.sqrt( J *( J + 1 ) ) )
psis   = np.linspace( 0.,  2.* np.pi, Mpoints )  # Random orientation of figure axis around J 

#X = np.zeros( (Npoints, Mpoints) )
#Y = np.zeros( (Npoints, Mpoints) )
#Z = np.zeros( (Npoints, Mpoints) )

#plt.ion()

# Also calculate Beta distributions
BetaValues           = np.linspace(0.0, 90.0, num=181, endpoint=True)
DeltaBeta            = BetaValues[1] - BetaValues[0]
Sigma                = 2.0 # degree
BetaDistribution     = np.zeros( len( BetaValues ))
BetaHistogram        = np.zeros( len( BetaValues ))


art3d.zalpha = lambda *args:args[0]

fig = plt.figure(figsize=(10,8))

ax  = fig.add_subplot( 221, projection='3d', aspect='equal')


X = []
Y = []
Z = []

for i in range(len(gammas)):
	for j in range(len(psis)):

		gamma = gammas[i]
		psi   = psis[j]

		RotationalMatrix1 = CalculateRotationalMatrix( alpha, beta, gamma )
		RotationalMatrix2 = CalculateRotationalMatrix( phi, theta, psi )

		#  Compose the two rotations in one Rotational Matrix and print in in file
		RotationalMatrix = np.dot( RotationalMatrix2, RotationalMatrix1 ) 

		CHRotated = np.dot( RotationalMatrix, CH )

		#X[i][j] = CHRotated[0]
		#Y[i][j] = CHRotated[1]
		#Z[i][j] = CHRotated[2]
		X.append( CHRotated[0] )
		Y.append( CHRotated[1] )
		Z.append( CHRotated[2] )

        # Only down to 90 degress (distribution is symmetric)
        #Beta = np.degrees( np.arccos( CHRotated[2] )) 
		Beta = np.degrees( np.arccos( CHRotated[2] )) 
		for k in range( len(BetaValues)):
			BetaDistribution[k] = BetaDistribution[k] + np.exp( -( BetaValues[k] - Beta  )**2. / ( 2. * Sigma **2. ))
			#if Beta >= BetaValues[k] and Beta < BetaValues[k] + DeltaBeta :
			#	BetaHistogram[k] = BetaHistogram[k] + 1.
			#	break
		#ax.scatter( CHRotated[0], CHRotated[1], CHRotated[2] )


#plt.show()
#quit()

X = np.array( X ) 
Y = np.array( Y )
Z = np.array( Z )

Density = stats.gaussian_kde( np.vstack([X,Y,Z]) )( np.vstack([X,Y,Z]) )

Density = Density / sum(Density)

#print min(Density), max(Density)
index   = Density.argsort()

X,Y,Z,Density = X[index],Y[index],Z[index],Density[index]
#ax.plot_surface( X,Y,Z)
#ax.plot_trisurf( X,Y,Z, cmap=Density)
cax = ax.scatter(X,Y,Z, s=20, edgecolor='',  c=Density , vmin=0, vmax=0.0003753 )
#cax.set_edgecolors = cax.set_facecolors = lambda *args:None
ax.set_title( "M = 1" )

# Also generate points to complete the sphere
XSphere = []
YSphere = []
ZSphere = []

phis   = np.linspace( 0.,  2.* np.pi, Npoints )
costhetas = np.linspace( -1.,  1., Npoints )
for i in range(len(phis)):
	for j in range(len(costhetas)):
		RotationalMatrix = CalculateRotationalMatrix( 0. , np.arccos(costhetas[j]) , phis[i] )

		CHRotated = np.dot( RotationalMatrix, CH )
		if CHRotated[2] > max( Z ) or CHRotated[2] < min( Z ):
			XSphere.append( CHRotated[0] )
			YSphere.append( CHRotated[1] )
			ZSphere.append( CHRotated[2] )

XSphere = np.array( XSphere )
YSphere = np.array( YSphere )
ZSphere = np.array( ZSphere )


#print len( ZSphere )
#print XSphere, YSphere, ZSphere
DensitySphere = stats.gaussian_kde( np.vstack([XSphere,YSphere,ZSphere]) )( np.vstack([XSphere,YSphere,ZSphere]) ) 

index   = DensitySphere.argsort()

XSphere,YSphere,ZSphere,DensitySphere = XSphere[index],YSphere[index],ZSphere[index],DensitySphere[index]

for i in range(len( DensitySphere)):
	DensitySphere[i] = 0.

cax = ax.scatter(XSphere,YSphere,ZSphere, s=20, edgecolor='', c=DensitySphere, vmin=0, vmax=0.0003753 )
#cax.set_edgecolors = cax.set_facecolors = lambda *args:None
ax.set_xlim( [-1.0, 1.0 ])
ax.set_ylim( [-1.0, 1.0 ])
ax.set_zlim( [-1.0, 1.0 ])
ax.set_axis_off()

# Draw Z axis 
ax.plot( [0., 0.], [0., 0.], [-1.3, -1.0], color='black', ls='-', alpha=0.8, lw=2)
ax.plot( [0., 0.], [0., 0.], [-1.0,  1.0], color='black', ls='--', alpha=0.8, lw=2)
ax.plot( [0., 0.], [0., 0.], [ 1.0,  1.3], color='black', ls='-', alpha=0.8, lw=2)

# Set initial view
ax.view_init( 10. )


ax  = fig.add_subplot(223) #, aspect='equal')
for k in range( len(BetaValues)):
	if BetaValues[k] < 45. :
		BetaDistribution[k] = 0.
Distribution = BetaDistribution/sum( BetaDistribution )/DeltaBeta
#Distribution = BetaHistogram/sum( BetaHistogram ) /DeltaBeta
#for k in range( len(BetaValues)):
MaxBetaDistribution = max( Distribution )
ax.set_xlim( [0, 90.0 ])
ax.yaxis.set_ticks( [ ] )
ax.xaxis.set_ticks( np.arange( 0., 90. +22.5 , 22.5) )
ax.set_xlabel( r'$\beta$ / $\degree$', labelpad=15 )
#ax.set_ylim( [0, 0.5 ])
ax.set_ylabel( 'Distribution / arb. ' )
ax.plot( BetaValues, Distribution, 'r-' )
fig.subplots_adjust(bottom=0.15)

############################################################################
# Add M = 0 sphere: plot ( 2, 1 )

# J,K,M value
J = 1.
K = 0.
M = 0.

# beta and theta are fixed to the new values
beta   = np.arccos(  K / np.sqrt( J *( J + 1 ) ) )
theta = np.arccos( M  / np.sqrt( J *( J + 1 ) ) )

# Create subplots
ax  = fig.add_subplot( 222, projection='3d', aspect='equal')

X = []
Y = []
Z = []

BetaDistribution     = np.zeros( len( BetaValues ))

for i in range(len(gammas)):
    for j in range(len(psis)):

        gamma = gammas[i]
        psi   = psis[j]

        RotationalMatrix1 = CalculateRotationalMatrix( alpha, beta, gamma )
        RotationalMatrix2 = CalculateRotationalMatrix( phi, theta, psi )

        #  Compose the two rotations in one Rotational Matrix and print in in file
        RotationalMatrix = np.dot( RotationalMatrix2, RotationalMatrix1 )

        CHRotated = np.dot( RotationalMatrix, CH )

        #X[i][j] = CHRotated[0]
        #Y[i][j] = CHRotated[1]
        #Z[i][j] = CHRotated[2]
        X.append( CHRotated[0] )
        Y.append( CHRotated[1] )
        Z.append( CHRotated[2] )
	
		# Do not calculate Beta values: distribution in beta is constant 
        #Beta = np.degrees( np.arccos( abs(CHRotated[2] )))
        #for k in range( len(BetaValues)):
        #	BetaDistribution[k] = BetaDistribution[k] + np.exp( -( BetaValues[k] - Beta  )**2. / ( 2. * Sigma **2. ))

X = np.array( X )
Y = np.array( Y )
Z = np.array( Z )

Density = stats.gaussian_kde( np.vstack([X,Y,Z]) )( np.vstack([X,Y,Z]) )

Density = Density / sum(Density)
index   = Density.argsort()

X,Y,Z,Density = X[index],Y[index],Z[index],Density[index]

cax = ax.scatter(X,Y,Z, s=20, edgecolor='',  c=Density, vmin=0, vmax=0.0003753 )

ax.set_title( "M = 0" )

ax.set_xlim( [-1.0, 1.0 ])
ax.set_ylim( [-1.0, 1.0 ])
ax.set_zlim( [-1.0, 1.0 ])

ax.set_axis_off()

# Draw Z axis 
ax.plot( [0., 0.], [0., 0.], [-1.3, -1.0], color='black', ls='-', alpha=0.8, lw=2)
ax.plot( [0., 0.], [0., 0.], [-1.0,  1.0], color='black', ls='--', alpha=0.8, lw=2)
ax.plot( [0., 0.], [0., 0.], [ 1.0,  1.3], color='black', ls='-', alpha=0.8, lw=2)

# Set initial view
ax.view_init( 10. )


############################################################################
# Add M = 0 distribution: plot ( 2, 2 )

ax  = fig.add_subplot(224) #, aspect='equal')

# The distribution is simply a constant
for k in range( len(BetaValues)):
	BetaDistribution[k] = 1.0

# Normalize 
Distribution = BetaDistribution/sum( BetaDistribution )/DeltaBeta

# Set X axis
ax.set_xlim( [0, 90.0 ])
ax.xaxis.set_ticks( np.arange( 0., 90. + 22.5 , 22.5) )
ax.set_xlabel( r'$\beta$ / $\degree$', labelpad=15 )

# Set Y axis
ax.set_ylim( [0., MaxBetaDistribution ])
ax.yaxis.set_ticks( [ ])
ax.set_ylabel( 'Distribution / arb.' )

# Plot
ax.plot( BetaValues, Distribution, 'r-' )

# Leave some space from bottom
fig.subplots_adjust(bottom=0.15)

# Leave some space on rightside for colour bar
fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.55, 0.05, 0.35 ])
fig.colorbar( cax, cax=cbar_ax  )

# Add letters 
fig.text( 0.15, 0.85, "A" , size="x-large" )
fig.text( 0.75, 0.85, "B" , size="x-large" )
fig.text( 0.15, 0.17, "C" , size="x-large" )
fig.text( 0.75, 0.17, "D" , size="x-large" )

# Done! 
plt.show()