#!/usr/bin/env python
import sys, commands
from os import path, getcwd, popen
from os import chdir, system
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from scipy.interpolate import spline

folders = ['74.0LO', '87.8LO', '102.9LO', '119.5LO', '102.9LO_475K', '119.5LO_475K']
#folders = ['74.0LO', '87.8LO', '102.9LO', '119.5LO']

angles = []
angles_scat = []
angles_trap = []
here=getcwd()
def collectangles(folder):
	chdir(folder)

	Ntraj = int( popen('ls | egrep "^[0-9]...$" | wc -l').read() )
	here_tmp=getcwd()				#workdir
	for i in range(1,(Ntraj+1)):
		print i
		traj= '%.4d' % i		# define working directory
		newdir=path.join(here_tmp, traj)	# updating mypath

		print newdir
		if not path.exists(newdir): 
			print "folder", newdir, "not found. exiting..."
			quit()

		chdir(newdir)

		if not path.exists('outcome') or not path.exists('PostAnalysis.dat'):
			continue

		outcome = open('outcome', 'r').readline()
		post_dat = open('PostAnalysis.dat', 'r').readlines()
		if 'REACTION ' in outcome:
#			print 'reaction'
			angles.append( [float(post_dat[25].split()[0]), float(post_dat[25].split()[1]), float(post_dat[25].split()[2]), float(post_dat[27].split()[0]), float(post_dat[27].split()[1]), float(post_dat[29].split()[0]), float(post_dat[29].split()[1]),
					float(post_dat[25].split()[4]), float(post_dat[25].split()[5]), float(post_dat[25].split()[6]), float(post_dat[27].split()[3]), float(post_dat[27].split()[4]), float(post_dat[29].split()[3]), float(post_dat[29].split()[4])] )
		elif 'SCATTERING' in outcome:
#			print 'scattering'
			angles_scat.append( [float(post_dat[24].split()[0]), float(post_dat[24].split()[1]), float(post_dat[24].split()[2]), float(post_dat[26].split()[0]), float(post_dat[26].split()[1]), float(post_dat[28].split()[0]), float(post_dat[28].split()[1])] )
		elif 'TRAPPING' in outcome:
#			print 'trapping'
			angles_trap.append( [float(post_dat[24].split()[0]), float(post_dat[24].split()[1]), float(post_dat[24].split()[2]), float(post_dat[26].split()[0]), float(post_dat[26].split()[1]), float(post_dat[28].split()[0]), float(post_dat[28].split()[1])] )

	chdir(here)

xlabels = [r'b ($\theta_1$)', r'c ($\theta_{2,3}$)', r'd ($\theta_3$)', r"e ($\beta'$)", r'd ($\beta$)', r'$\gamma_1$', r'$\gamma$']

for i in range(len(folders)):
	collectangles(folders[i])

angles_theta23 = []
angles_theta23_ts = []
angles_theta23_scat = []
angles_theta23_trap = []
for i in range(len(angles)):
	angles_theta23.append( angles[i][1] )
	angles_theta23.append( angles[i][2] )
	angles_theta23_ts.append( angles[i][8] )
	angles_theta23_ts.append( angles[i][9] )
for i in range(len(angles_scat)):
	angles_theta23_scat.append( angles_scat[i][1] )
	angles_theta23_scat.append( angles_scat[i][2] )
for i in range(len(angles_trap)):
	angles_theta23_trap.append( angles_trap[i][1] )
	angles_theta23_trap.append( angles_trap[i][2] )

print 'Theta2,3:'
print np.mean(angles_theta23), np.std(angles_theta23, ddof=1), np.std(angles_theta23, ddof=1) / len(angles_theta23)**0.5
print np.mean(angles_theta23_ts), np.std(angles_theta23_ts, ddof=1), np.std(angles_theta23_ts, ddof=1) / len(angles_theta23_ts)**0.5
print np.mean(angles_theta23_scat), np.std(angles_theta23_scat, ddof=1), np.std(angles_theta23_scat, ddof=1) / len(angles_theta23_scat)**0.5
print np.mean(angles_theta23_trap), np.std(angles_theta23_trap, ddof=1), np.std(angles_theta23_trap, ddof=1) / len(angles_theta23_trap)**0.5

if len(angles) != 0 and len(angles_scat) != 0 and len(angles_trap) != 0:
	N = len( angles )
	mu = []
	sigma = []
	error = []
	angles = np.array(angles)
	angles_scat = np.array(angles_scat)
	angles_trap = np.array(angles_trap)
	for i in range( len(xlabels) * 2):
		mu.append( np.mean(angles[:,i]) )
		sigma.append( np.std(angles[:,i], ddof=1) )
		error.append( sigma[i] / len(angles)**0.5 )
	for i in range( len(xlabels) ):
		mu.append( np.mean(angles_scat[:,i]) )
		sigma.append( np.std(angles_scat[:,i], ddof=1) )
		error.append( sigma[i+len(xlabels)*2] / len(angles_scat)**0.5 )
	for i in range( len(xlabels) ):
		mu.append( np.mean(angles_trap[:,i]) )
		sigma.append( np.std(angles_trap[:,i], ddof=1) )
		error.append( sigma[i+len(xlabels)*3] / len(angles_trap)**0.5 )
	
	print 'Reacted (t=0):'
	for i in range( len(xlabels) ):
		print '{:s}: {:.1f} +/- {:.1f} ({:.1f})'.format( xlabels[i], mu[i], error[i], sigma[i] )

	print 'Reacted (r=r_TS):'
	for i in range( len(xlabels) ):
		print '{:s}: {:.1f} +/- {:.1f} ({:.1f})'.format( xlabels[i], mu[i+len(xlabels)], error[i+len(xlabels)], sigma[i+len(xlabels)] )

	print 'Scattered (t=0):'
	for i in range( len(xlabels) ):
		print '{:s}: {:.1f} +/- {:.1f} ({:.1f})'.format( xlabels[i], mu[i+len(xlabels)*2], error[i+len(xlabels)*2], sigma[i+len(xlabels)*2] )

	print 'Trapped (t=0):'
	for i in range( len(xlabels) ):
		print '{:s}: {:.1f} +/- {:.1f} ({:.1f})'.format( xlabels[i], mu[i+len(xlabels)*3], error[i++len(xlabels)*3], sigma[i++len(xlabels)*3] )

output = open('angles_reac.dat', 'w')
for i in range( len(angles) ):
	for j in range(6):
		output.write( str( angles[i,j] ) )
		output.write('\t')
	output.write('\n')

output = open('angles_scat.dat', 'w')
for i in range( len(angles_scat) ):
	for j in range(3):
		output.write( str( angles_scat[i,j] ) )
		output.write('\t')
	output.write('\n')

output = open('angles_trap.dat', 'w')
for i in range( len(angles_trap) ):
	for j in range(3):
		output.write( str( angles_trap[i,j] ) )
		output.write('\t')
	output.write('\n')

nrows=5
ncols=1

fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(3.37,7))
#mpl.rc("font", size=10)
mpl.rc("text", usetex=True)

plt.subplot(nrows,ncols,1)
angles_CH4 = []
dat = open('ni_angles_reac.dat', 'r').readlines()
for i in range(1, len(dat)):
	angles_CH4.append( [float(dat[i].split()[0]), float(dat[i].split()[3])] )
angles_CH4 = np.array(angles_CH4)

n, bins = np.histogram( angles_CH4[:,0], bins=20, density=True )
x = []
x.append( bins[0] - (bins[1]-bins[0])/2 )
for i in range(len(bins)-1):
	x.append( (bins[i]+bins[i+1])/2. )
x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
n = np.insert(n, 0, 0)
n = np.append(n,0)
xnew = np.linspace(min(x),max(x),200)
n_smooth = spline(x,n,xnew)
plt.plot(xnew,n_smooth, linestyle='--', color='purple')

n, bins = np.histogram( angles_CH4[:,1], bins=20, density=True )
x = []
x.append( bins[0] - (bins[1]-bins[0])/2 )
for i in range(len(bins)-1):
	x.append( (bins[i]+bins[i+1])/2. )
x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
n = np.insert(n, 0, 0)
n = np.append(n,0)
xnew = np.linspace(min(x),max(x),200)
n_smooth = spline(x,n,xnew)
plt.plot(xnew,n_smooth, label=r'CHD$_3$', color='purple')

angles_HOD = []
dat = open('HODNi111.dat', 'r').readlines()
for i in range(1, len(dat)):
	if dat[i].split()[4] == 'REAC':
		angles_HOD.append( [float(dat[i].split()[10]), float(dat[i].split()[11])] )
angles_HOD = np.array(angles_HOD)

n, bins = np.histogram( angles_HOD[:,0], bins=20, density=True )
x = []
x.append( bins[0] - (bins[1]-bins[0])/2 )
for i in range(len(bins)-1):
	x.append( (bins[i]+bins[i+1])/2. )
x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
n = np.insert(n, 0, 0)
n = np.append(n,0)
xnew = np.linspace(min(x),max(x),200)
n_smooth = spline(x,n,xnew)
plt.plot(xnew,n_smooth, linestyle='--', color='dimgray')

n, bins = np.histogram( angles_HOD[:,1], bins=20, density=True )
x = []
x.append( bins[0] - (bins[1]-bins[0])/2 )
for i in range(len(bins)-1):
	x.append( (bins[i]+bins[i+1])/2. )
x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
n = np.insert(n, 0, 0)
n = np.append(n,0)
xnew = np.linspace(min(x),max(x),200)
n_smooth = spline(x,n,xnew)
plt.plot(xnew,n_smooth, label='HOD', color='dimgray')

x = np.linspace(0.,180.,180.)
n = np.sin( np.radians(x) )
n /= sum(n)
plt.plot(x,n, label='Statistical', color='k')

y_limits = plt.ylim()
plt.plot( [135.7, 135.7], plt.ylim(), linestyle='dotted', c='purple' )
plt.plot( [123., 123.], plt.ylim(), linestyle='dotted', c='dimgray' )
plt.ylim(0., y_limits[1])

plt.tick_params(labelleft='off')
plt.xlim(0,180)
plt.legend(loc='center left', fontsize=8, frameon=False)
plt.tick_params(labelbottom='off')
plt.annotate(r'a ($\theta_d$)', xy=(15, y_limits[1]*0.8), size=12)

angle_list = [1, 2, 5, 4]

for plot_idx in range(2,nrows+1):
	plt.subplot(nrows,ncols,plot_idx)

	idx = angle_list[plot_idx-2]
	if len(angles) != 0:
		if plot_idx == 3:
			n, bins = np.histogram( angles_theta23[:], bins=20, density=True )
		else:
			n, bins = np.histogram( angles[:,idx-1], bins=20, density=True )
		x = []
		x.append( bins[0] - (bins[1]-bins[0])/2 )
		for i in range(len(bins)-1):
			x.append( (bins[i]+bins[i+1])/2. )
		x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
		n = np.insert(n, 0, 0)
		n = np.append(n,0)

		xnew = np.linspace(0.,180.,180.)
		n_smooth = spline(x,n,xnew)
		plt.plot(xnew,n_smooth, label='Reacted ($t = 0$ fs)')

		if plot_idx == 3:
			n, bins = np.histogram( angles_theta23_ts[:], bins=20, density=True )
		else:
			n, bins = np.histogram( angles[:,idx+6], bins=20, density=True )
		x = []
		x.append( bins[0] - (bins[1]-bins[0])/2 )
		for i in range(len(bins)-1):
			x.append( (bins[i]+bins[i+1])/2. )
		x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
		n = np.insert(n, 0, 0)
		n = np.append(n,0)

		xnew = np.linspace(0.,180.,180.)
		n_smooth = spline(x,n,xnew)
		plt.plot(xnew,n_smooth, label=r'Reacted ($r = 1.6$ \r{A})')

	if len(angles_scat) != 0:
		angles_scat = np.array(angles_scat)
		if plot_idx == 3:
			n, bins = np.histogram( angles_theta23_scat[:], bins=20, density=True )
		else:
			n, bins = np.histogram( angles_scat[:,idx-1], bins=50, density=True )
		x = []
		x.append( bins[0] - (bins[1]-bins[0])/2 )
		for i in range(len(bins)-1):
			x.append( (bins[i]+bins[i+1])/2. )
		x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
		n = np.insert(n, 0, 0)
		n = np.append(n,0)

		xnew = np.linspace(0.,180.,180.)
		n_smooth = spline(x,n,xnew)
		plt.plot(xnew,n_smooth, label='Scattered ($t = 0$ fs)')

	if len(angles_trap) != 0:
		angles_trap = np.array(angles_trap)
		if plot_idx == 3:
			n, bins = np.histogram( angles_theta23_trap[:], bins=20, density=True )
		else:
			n, bins = np.histogram( angles_trap[:,idx-1], bins=20, density=True )
		x = []
		x.append( bins[0] - (bins[1]-bins[0])/2 )
		for i in range(len(bins)-1):
			x.append( (bins[i]+bins[i+1])/2. )
		x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
		n = np.insert(n, 0, 0)
		n = np.append(n,0)

		xnew = np.linspace(0.,180.,180.)
		n_smooth = spline(x,n,xnew)
		plt.plot(xnew,n_smooth, label='Trapped ($t = 0$ fs)', c='dimgray')

	if idx == 1: TS_angle = 115.6
	elif idx == 2: TS_angle = 55.2
	elif idx == 3: TS_angle = 55.3
	elif idx == 4: TS_angle = 108.4
	elif idx == 5: TS_angle = 177.2
	elif idx == 6: TS_angle = 135.9
	elif idx == 7: TS_angle = 61.6

	y_limits = plt.ylim()
	plt.plot( [TS_angle, TS_angle], plt.ylim(), linestyle='dotted', c='k' )
	plt.ylim(0., y_limits[1])

	plt.tick_params(labelleft='off')
	plt.xlim(0,180)
	plt.annotate(xlabels[idx-1], xy=(15, y_limits[1]*0.8), size=12)
#	if plot_idx == 1:
#		plt.annotate(xlabels[plot_idx-1], xy=(160, y_limits[1]*0.8), size=12)
#	if plot_idx == 2:
#		plt.annotate(xlabels[plot_idx-1], xy=(160, y_limits[1]*0.8), size=12)
#	if plot_idx == 3:
#		plt.annotate(xlabels[plot_idx-1], xy=(160, y_limits[1]*0.8), size=12)
	if plot_idx == 4:
		plt.legend(loc='upper center', fontsize=8, frameon=False)
	if plot_idx != nrows:
		plt.tick_params(labelbottom='off')
	if plot_idx == 3:
		plt.ylabel('Distribution (arb.)')
	if plot_idx == 5:
		plt.xlabel('Angle (deg)')

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