#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from random import choice, random
from scipy import interpolate

Nrows=4
Ncols=2
#fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(6.69,8.))
plt.figure(1)
plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(6.69,8.))
mpl.rc("text", usetex=True)

def ptheta( j, mj, theta, direction, atomicunits=True ):
	'''
	|L_theta| = hbar * [ j(j+1) - (mJ / sin(theta))^2 ]^0.5
	|L_phi| = hbar * mJ
	Note that the momentum has a direction, i.e., positive or negative
	'''
	if atomicunits:
		unitconv = 1.
	else:
		unitconv = 0.529177249 / 1822.888486209 * 21.876913	# angstrom to bohr, electron mass to amu, velocity in Hartree atomic units to angstrom / fs
	if mj == 0:
		pt = unitconv * direction * np.sqrt( j*(j+1) )
	else:
		pt = unitconv * direction * np.sqrt( j*(j+1) - ( mj / np.sin(theta) )**2 )
	
	return pt

def rotperiod( j, mu, r0 ):
	'''
	I=mu*R**2
	omega=2*pi*frequency
	L=I*omega
	T=1/frequency=2*pi*mu*R**2/L
	'''
	L = abs( ptheta(j, 0, 0, 1, False) )
	period = 2.*np.pi*mu*r0**2 / L

	return period

def theta_( j, mj, phase ):
	cos_beta = 0.
	if mj != 0:
		cos_beta = mj / np.sqrt( j*(j+1) )
	BETA = np.arccos( cos_beta )
	sin_beta = np.sin( BETA )

	gamma = 2.*np.pi*random()
	cos_gamma = np.cos( gamma )
	cos_theta = -sin_beta * cos_gamma
	THETA_i = np.arccos( cos_theta )
	if gamma < np.pi:
		PT_i = ptheta(j, mj, THETA_i, 1, False)
	else:
		PT_i = ptheta(j, mj, THETA_i, -1, False)

	gamma = ( gamma + phase ) % ( 2.*np.pi )
	cos_gamma = np.cos( gamma )
	cos_theta = -sin_beta * cos_gamma
	THETA_f = np.arccos( cos_theta )
	if gamma < np.pi:
		PT_f = ptheta(j, mj, THETA_f, 1, False)
	else:
		PT_f = ptheta(j, mj, THETA_f, -1, False)
	
	return THETA_i, THETA_f, PT_i, PT_f

#****** V E L O C I T Y   D I S T R I B U T I O N *************************************************
#
#      F. Nattino, 05/2011
#
#     Modified to python by N. Gerrits, 01/2019
#
#     Select incidence energy according to flux weighted velocity distribution as 
#      in Michelsen and al. J Chem. Phys. 94, 7502 (1991) 
#    
#     G(E)dE = ( 1 / N ) * E * exp[ - 4Es * ( SQRT(E) - SQRT(Es) ) ** 2 / DelEs**2 ] dE
#
#     DelEs / Es = 2 DelVs / Vs = 2 / S
#
#     N = (DelEs / 2 S)**2 * { ( 1 + S**2 ) * exp( -S**2) + 
#         SQRT(PI) * S * ( 1.5 + S**2 )[ 1 + ERF(S) ] }
#
#     Units: Mass = kg, Vs and DelVs in m/sec, Etransmax in eV
#
#     USE: Constants (J2eV and ev2kjmol) and Random

def SetupVelocityDistribution():
	from math import erf
	J2eV=6.24181 * 10**18
	ev2kjmol=96.4853365
	amu2kg=1.66053892 * 10**-27
	class vel:
		"""Translational information"""
		pass
	# Translational parameters
	vel.dist	= True	# Use a Velocity distribution for normal translational energy
	vel.stream	= 3616. # Stream Velocity (m/s), if vel.dist is false then this is used as the initial translational energy
	vel.width	= 371.	# Velocity width (m/s)
	vel.Emax	= 4.5	# Maximum translational energy considered (eV)
	vel.shift	= [False, 0.0]	# Shift velocity distribution. If true, by how much (eV)
	vel.Etrans0	= 0.	# Translational energy (eV) if no velocity distribution is used, i.e. a monocromatic beam
	vel.output	= False	# If True, the velocity distribution will be written to a file (VelocityDistribution.dat) and a summary to the screen,  and a normalization check is performed
	vel.mass	= 36.458# Mass of HCl (amu)

	vel.Estream = 0.5 * vel.mass*amu2kg * vel.stream**2. * J2eV
	vel.DelEstream = 2. * vel.Estream * vel.width / vel.stream
	vel.SqEstream = np.sqrt( vel.Estream )

	srma = vel.stream / vel.width
	# Normalization factor
	vel.FacNor = ( vel.DelEstream /( 2. * srma) )**2. * (( 1. + srma**2. )* np.exp( -1. * srma**2. ) + np.sqrt( np.pi ) * srma * ( 1.5 + srma**2. ) * ( 1. + erf( srma ) ) )
	EGmax = (( vel.DelEstream**2. ) + 2. * ( vel.Estream**2. ) + 2. * vel.Estream * np.sqrt( ( vel.Estream**2. ) + ( vel.DelEstream ** 2.)))/(4. * vel.Estream)
	Expon = 4.* vel.Estream * ( np.sqrt(EGmax) - vel.SqEstream )**2. / vel.DelEstream**2.
	Expfac = np.exp( -1.* Expon)

	vel.Gmax = EGmax / vel.FacNor * Expfac

	# Distribution is checked and an output is written
	if vel.output == True:
		Edist = open('VelocityDistribution.dat', 'w')
		Edist.write( '# Normalized Translational Energy Distribution. E in eV.' )
		StepIntegration = 100000
		DelE = vel.Emax / float( StepIntegration )
		Eaver = 0.
		Dnorm = 0.
		for i in range( StepIntegration ):
			E = float(i) * DelE
			Expon = 4. * vel.Estream * (np.sqrt( E ) - vel.SqEstream)**2. / vel.DelEstream **2.
			Expfac = np.exp( - Expon )
			GE = Expfac * E
			Eaver += GE * E * DelE
			Dnorm += GE * DelE
			Edist.write('{:f}	{:f}'.format( E, (GE / vel.FacNor) ) )

		Eaver = Eaver / vel.FacNor
		Dnorm = Dnorm / vel.FacNor

		print( ' Energy corresponding to Stream Velocity:   E0 (eV) = ', vel.Estream )
		print( ' DeltaE (eV) = ', vel.DelEstream )
		print( ' Mass (amu) = ', vel.mass )
		print( ' Energy at Maximum Distribution (eV) = ', EGmax )
		print( ' The average collision energy is (eV) = '    ,  Eaver )
		print( ' The average collision energy is (kJ/mol) = ' ,  Eaver * ev2kjmol )
		print( ' Verify normalization of the distribution: {:12.10f}\n'.format( Dnorm ) )

		if (abs( Dnorm - 1. ) > 1.E-4 ):
			print( 'WARNING: Increase Etransmax' )
			print( ' Normalization of the distribution: ',  Dnorm )
			exit(1)

	return vel

def SelectVelocityFromDistribution( vel ):
	J2eV=6.24181 * 10**18
	amu2kg=1.66053892 * 10**-27
	Gran = 99999.
	G = 0.
	while Gran > G:
		Etrans = np.random.random_sample() * vel.Emax
		Expon = 4.* vel.Estream * ( np.sqrt( Etrans ) - vel.SqEstream )**2. / vel.DelEstream**2.
		Expfac = np.exp( -1.* Expon)
		G = Etrans / vel.FacNor * Expfac
		Gran = np.random.random_sample() * vel.Gmax

	VelTrans = np.sqrt( 2. * Etrans / J2eV / ( vel.mass * amu2kg) ) * 1.E-5

	return VelTrans

def theta_final( j, mj, vel ):
	d = 5.5	# Distance between initial Z and Z_TS in Angstrom
	#v0 = 3500. * 10**-5	# 10**10 / 10**15, m/s to Angstrom/fs --> Velocity in Angstrom/fs
	if vel.dist == True:
		v0 = SelectVelocityFromDistribution( vel )
	else:
		v0 = vel.stream * 10**-5
	mu = 0.97	# Reduced mass in amu
	r0 = 1.27	# Equilibrium gas phase HCl bond length in Angstrom
	period = rotperiod( j, mu, r0 )
	phase = -( ( d / v0 / period ) % 1. ) * 2.*np.pi
	
	theta_i, theta_f, pt_i, pt_f = theta_( j, mj, phase )
	theta_i_DEG = np.degrees(theta_i)

	return [ theta_i_DEG, pt_i, theta_weight(theta_i_DEG), np.degrees(theta_f) ]

def theta_weight( THETA ):
	'''
	This was how trajectories were weighted in the original plots.
	This relies on the fact that the statistical initial distribution is a sin(theta) distribution.
	Note that the input expects the theta angle to be in degrees instead of radians.
	'''
	weight = 1. / np.sin( np.radians( THETA ) )
	#weight = 0.5 * np.exp( -(np.radians( THETA ) - 80.)**2 / 1000. ) + 2*np.exp( -(np.radians( THETA ) - 170.)**2 / 8000. ) - 1.	# A function derived by Jan Geweke to describe reactivity

	return weight

def theta_reactivity( THETA, nu ):
	'''
	This function yields the reactivity of a particular theta value determined approximately via the reactivity of v=0, J=0.
	Note that the input expects the theta angle to be in degrees instead of radians.
	'''
	if nu==0:
		theta_sticking = [0.0000, 0.0019, 0.0132, 0.1494, 0.3149, 0.3544, 0.3224, 0.3913, 0.3926, 0.4464, 0.4814, 0.4561]
	elif nu==1:
		theta_sticking = [0.0000, 0.0000, 0.0327, 0.1857, 0.4057, 0.4968, 0.4798, 0.5634, 0.7256, 0.7260, 0.7364, 0.7748]
	elif nu==2:
		theta_sticking = [0.0048, 0.0141, 0.0503, 0.2008, 0.4619, 0.5459, 0.5474, 0.6561, 0.8207, 0.8576, 0.8716, 0.8416]
	
	theta = np.arange(10,180,15)
	f = interpolate.UnivariateSpline(theta, theta_sticking, s=0)
	
	return f( THETA )

def plotangles_analytical(idx, title, nu, J, N, vel):
	ax2 = plt.subplot(Nrows,Ncols,idx)

	mJ = np.arange( -J, J+1, 1 )
	angles = []
	for i in range( N ):
		angles.append( theta_final( J, choice( mJ ), vel ) )
	angles = np.array( angles )

	angles_stat_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[min(angles[:,1]), max(angles[:,1])], [20., 160.]], bins=36, density=True, weights=angles[:,2] )
	angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[min(angles[:,1]), max(angles[:,1])], [20., 160.]], bins=36, density=True, weights=angles[:,2]*theta_reactivity(angles[:,3], nu) )

	angles_binned -= angles_stat_binned	# In order to get the relative distribution of reactive trajectories vs the initial distribution
	angles_binned += 1.			# Arbitrary shift so that if reacted/statistical=1 it also yields a probability density of 1
	
	vlimit = max( abs( np.amin(angles_binned) ), abs( np.amax(angles_binned) ) )
	plt.pcolormesh(yedges, xedges, angles_binned, vmin=1.-(vlimit-1.), vmax=vlimit, cmap='bwr')
	#plt.pcolormesh(yedges, xedges, angles_binned, cmap='jet')
	#plt.pcolormesh(yedges, xedges, angles_stat_binned, cmap='jet')
	#cbar = plt.colorbar(label='Relative intensity')
	#cbar = plt.colorbar(label='Relative intensity', ticks=[1.])
	#cbar.ax.set_yticklabels(['Avg.'])
	cbar = plt.colorbar(label='Relative intensity', ticks=[1.-(vlimit-1.),1.,vlimit])
	cbar.ax.set_yticklabels(['Low', 'Avg.', 'High'])

	plt.annotate(title, xy=(0.525, 0.075), xycoords='axes fraction', size=10, color='k', bbox=dict(boxstyle='round', fc='w'))

	plt.xlabel(r'$\theta_\textrm{i}$ angle (degrees)')
	plt.ylabel(r'$p_{\theta_\textrm{i}}$ (amu\r{A}$^2$/fs)')

	plt.xticks([20., 55., 90., 125., 160.])

	return

nu = 0
N = 100000
vel = SetupVelocityDistribution()	# Rather set up the velocity distribution only once
for J in range(1,9):
	plotangles_analytical(J, r'$\nu={0:d},J={1:d}$'.format(nu, J), nu, J, N, vel)

plt.tight_layout()
plt.savefig('momentum_theta_analytical_v{:d}.pdf'.format(nu))

#### From here we compute sticking as a function of initial theta

# J, S0 for theta 10, 30, ..., 170, error
S0_v0 = np.array(
[[0, 0.0000, 0.0019, 0.0132, 0.1494, 0.3149, 0.3544, 0.3224, 0.3913, 0.3926, 0.4464, 0.4814, 0.4561, 0.0000, 0.0019, 0.0040, 0.0112, 0.0131, 0.0132, 0.0129, 0.0141, 0.0159, 0.0182, 0.0249, 0.0466],
[2, 0.3333, 0.3773, 0.2808, 0.3733, 0.2962, 0.2363, 0.2861, 0.2859, 0.2693, 0.2901, 0.3945, 0.3223, 0.0348, 0.0191, 0.0182, 0.0134, 0.0128, 0.0125, 0.0134, 0.0125, 0.0139, 0.0178, 0.0234, 0.0425],
[4, 0.5320, 0.4973, 0.4565, 0.4700, 0.4513, 0.3648, 0.2913, 0.2104, 0.1478, 0.0745, 0.0188, 0.0000, 0.0350, 0.0211, 0.0170, 0.0175, 0.0127, 0.0139, 0.0133, 0.0113, 0.0114, 0.0102, 0.0062, 0.0000],
[6, 0.4468, 0.4037, 0.3911, 0.4099, 0.4374, 0.4344, 0.3788, 0.3240, 0.2309, 0.1517, 0.0954, 0.0455, 0.0363, 0.0202, 0.0174, 0.0150, 0.0137, 0.0142, 0.0138, 0.0136, 0.0136, 0.0126, 0.0163, 0.0199],
[8, 0.2019, 0.2032, 0.2257, 0.2565, 0.2746, 0.3245, 0.3855, 0.4247, 0.5000, 0.5467, 0.5871, 0.6379, 0.0275, 0.0169, 0.0145, 0.0136, 0.0122, 0.0135, 0.0137, 0.0147, 0.0162, 0.0185, 0.0246, 0.0446]]
)

S0_v1 = np.array(
[[0, 0.0000, 0.0000, 0.0327, 0.1857, 0.4057, 0.4968, 0.4798, 0.5634, 0.7256, 0.7260, 0.7364, 0.7748, 0.0000, 0.0000, 0.0062, 0.0123, 0.0141, 0.0140, 0.0141, 0.0147, 0.0149, 0.0169, 0.0224, 0.0396],
[2, 0.5562, 0.4976, 0.4092, 0.5189, 0.4213, 0.3741, 0.4264, 0.4481, 0.4173, 0.4340, 0.6329, 0.7168, 0.0372, 0.0200, 0.0203, 0.0140, 0.0142, 0.0144, 0.0149, 0.0138, 0.0157, 0.0198, 0.0234, 0.0424],
[4, 0.7111, 0.7249, 0.7202, 0.7058, 0.6551, 0.5141, 0.4157, 0.3274, 0.2482, 0.1312, 0.0605, 0.0602, 0.0338, 0.0195, 0.0157, 0.0164, 0.0123, 0.0146, 0.0145, 0.0131, 0.0140, 0.0133, 0.0111, 0.0261],
[6, 0.7088, 0.6837, 0.6974, 0.7261, 0.6604, 0.6288, 0.5178, 0.3798, 0.2778, 0.2101, 0.1957, 0.1518, 0.0337, 0.0195, 0.0166, 0.0138, 0.0133, 0.0141, 0.0144, 0.0143, 0.0146, 0.0143, 0.0221, 0.0339],
[8, 0.5023, 0.5486, 0.5716, 0.5065, 0.4744, 0.5098, 0.5582, 0.5909, 0.6397, 0.6839, 0.6572, 0.6842, 0.0343, 0.0211, 0.0172, 0.0158, 0.0138, 0.0146, 0.0143, 0.0149, 0.0159, 0.0176, 0.0241, 0.0435]]
)

S0_v2 = np.array(
[[0, 0.0048, 0.0141, 0.0503, 0.2008, 0.4619, 0.5459, 0.5474, 0.6561, 0.8207, 0.8576, 0.8716, 0.8416, 0.0048, 0.0053, 0.0077, 0.0128, 0.0142, 0.0141, 0.0144, 0.0147, 0.0134, 0.0137, 0.0175, 0.0363],
[2, 0.5511, 0.5353, 0.4468, 0.5759, 0.4619, 0.4276, 0.4799, 0.4940, 0.4387, 0.5000, 0.6959, 0.8100, 0.0375, 0.0200, 0.0211, 0.0142, 0.0146, 0.0153, 0.0153, 0.0141, 0.0159, 0.0203, 0.0234, 0.0392],
[4, 0.8708, 0.8110, 0.8256, 0.7961, 0.7402, 0.5667, 0.4737, 0.3738, 0.3257, 0.2436, 0.1290, 0.0759, 0.0251, 0.0177, 0.0136, 0.0146, 0.0116, 0.0148, 0.0150, 0.0137, 0.0154, 0.0171, 0.0159, 0.0298],
[6, 0.8580, 0.7877, 0.7964, 0.7885, 0.7738, 0.7345, 0.6076, 0.4787, 0.3220, 0.2507, 0.2466, 0.2400, 0.0263, 0.0176, 0.0150, 0.0131, 0.0122, 0.0132, 0.0144, 0.0149, 0.0155, 0.0157, 0.0251, 0.0427],
[8, 0.7551, 0.7685, 0.7402, 0.7331, 0.6331, 0.6016, 0.6152, 0.6482, 0.7130, 0.6816, 0.7180, 0.8051, 0.0307, 0.0184, 0.0156, 0.0145, 0.0136, 0.0148, 0.0144, 0.0146, 0.0153, 0.0177, 0.0230, 0.0365]]
)

theta_MD = np.arange(10,180,15)

def plot_sticking_theta(idx, title, nu, J, N, vel):
	mJ = np.arange( -J, J+1, 1 )
	angles = []
	for i in range( N ):
		angles.append( theta_final( J, choice( mJ ), vel ) )
	angles = np.array( angles )

	angles_stat_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[min(angles[:,1]), max(angles[:,1])], [5., 175.]], bins=170, density=False )
	angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[min(angles[:,1]), max(angles[:,1])], [5., 175.]], bins=170, density=False, weights=theta_reactivity(angles[:,3], nu) )

	if nu==0:
		plt.errorbar( theta_MD, S0_v0[int(J/2),1:len(theta_MD)+1], S0_v0[int(J/2),len(theta_MD)+1:], label=r'$J={:d}$'.format(J), capsize=4, color='C{:d}'.format( int(J/2) ), ls='' )
	elif nu==1:
		plt.errorbar( theta_MD, S0_v1[int(J/2),1:len(theta_MD)+1], S0_v1[int(J/2),len(theta_MD)+1:], label=r'$J={:d}$'.format(J), capsize=4, color='C{:d}'.format( int(J/2) ), ls='' )
	elif nu==2:
		plt.errorbar( theta_MD, S0_v2[int(J/2),1:len(theta_MD)+1], S0_v2[int(J/2),len(theta_MD)+1:], label=r'$J={:d}$'.format(J), capsize=4, color='C{:d}'.format( int(J/2) ), ls='' )
	
	sticking = []
	theta = []
	for i in range( len(angles_binned) ):
		sticking.append( sum(angles_binned[:,i]) / sum(angles_stat_binned[:,i]) )
		theta.append( (yedges[i] + yedges[i+1])/2. )
	plt.plot( theta, sticking, color='C{:d}'.format( int(J/2) ) )
	
	plt.annotate(title, xy=(0.5, 0.1), xycoords='axes fraction', horizontalalignment='center', size=10, color='k')

	if idx == 1:
		plt.legend(loc='upper center', numpoints=1, frameon=True, ncol=2, columnspacing=1., fontsize=8)
	plt.ylabel(r'Reaction probability')

	plt.xticks([0., 45., 90., 135., 180.])
	plt.xlim(0.,180.)
	plt.ylim(0., 0.9)
	
	if idx < 3:
		plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
	else:
		plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=True)
		plt.xlabel(r'$\theta_\textrm{i}$ angle (degrees)')

	return

plt.figure(2)
Nrows=3
Ncols=1
plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(3.69,5.))

N = 1000000
nu = 0
for nu in range(3):
	ax2 = plt.subplot(Nrows,Ncols,nu+1)
	for J in [2,4,6,8]:
		plot_sticking_theta(nu+1, r'$\nu={0:d}$'.format(nu), nu, J, N, vel)

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('sticking_theta_analytical.pdf')