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

Ntraj = [1000, 1000, 1000, 500, 500, 500, 1000, 1000, 500]
folders = [ '../CHD3_Pd111_LO/89.2LO', '../CHD3_Pd111_LO/102-5kJmol', '../CHD3_Pd111_LO/111-9kJmol', '../CHD3_Pd111_LO/120-0kJmol', '89.2v1', '97.4v1', '102.4v1', '111.8v1', '120.1v1' ]

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

	here_tmp=getcwd()				#workdir
	for i in range(1,(Ntraj+1)):
		traj= '%.4d' % i		# define working directory
		newdir=path.join(here_tmp, traj)	# updating mypath

		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:
			angles.append( [float(post_dat[25].split()[0]), float(post_dat[27].split()[0]), float(post_dat[29].split()[0]),
					float(post_dat[25].split()[2]), float(post_dat[27].split()[2]), float(post_dat[29].split()[2])] )
		else:
			angles_scat.append( [float(post_dat[24].split()[0]), float(post_dat[26].split()[0]), float(post_dat[28].split()[0])] )

	chdir(here)

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

if len(angles) != 0 and len(angles_scat) != 0:
	N = len( angles )
	mu = []
	sigma = []
	error = []
	angles = np.array(angles)
	angles_scat = np.array(angles_scat)
	for i in range(6):
		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(3):
		mu.append( np.mean(angles_scat[:,i]) )
		sigma.append( np.std(angles_scat[:,i], ddof=1) )
		error.append( sigma[i+6] / len(angles_scat)**0.5 )
	
	print 'Reacted (t=0):'
	print "theta\t\tbeta\t\talpha\t\tphi\t\tgamma1\t\tgamma2"
	print( '{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})'.format(mu[0], error[0], sigma[0], mu[1], error[1], sigma[1], mu[2], error[2], sigma[2]) )

	print 'Reacted (r=r_TS):'
	print "theta\t\tbeta\t\talpha\t\tphi\t\tgamma1\t\tgamma2"
	print( '{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})'.format(mu[3], error[3], sigma[3], mu[4], error[4], sigma[4], mu[5], error[5], sigma[5]) )

	print 'Scattered (t=0):'
	print "theta\t\tbeta\t\talpha\t\tphi\t\tgamma1\t\tgamma2"
	print( '{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})'.format(mu[6], error[6], sigma[6], mu[7], error[7], sigma[7], mu[8], error[8], sigma[8]) )

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')

fig, ax = plt.subplots(nrows=3, ncols=1, figsize=(3.37,4))
mpl.rc("font", size=10)

xlabels = [r'$\theta$', r'$\beta$', r'$\gamma$']

for plot_idx in range(1,4):
	plt.subplot(3,1,plot_idx)

	if len(angles) != 0:
		n, bins = np.histogram( angles[:,plot_idx-1], bins=10, 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)
		plt.plot(x,n, label='Reacted (t = 0 fs)')

		n, bins = np.histogram( angles[:,plot_idx+2], bins=10, 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)
		plt.plot(x,n, label='Reacted (r = 1.6 $\AA$)')

	if len(angles_scat) != 0:
		angles_scat = np.array(angles_scat)
		n, bins = np.histogram( angles_scat[:,plot_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)
		plt.plot(x,n, label='Scattered (t = 0 fs)')

	if plot_idx == 1: TS_angle = 135.9
	if plot_idx == 2: TS_angle = 165.0
	if plot_idx == 3: TS_angle =  29.1
	y_limits = plt.ylim()
	plt.plot( [TS_angle, TS_angle], plt.ylim(), linestyle='dotted', c='k' )
	plt.ylim(y_limits)

	plt.tick_params(labelleft='off')
	plt.xlim(0,180)
	if plot_idx == 1:
		plt.annotate(xlabels[plot_idx-1], xy=(40, 0.03), size=12)
	if plot_idx == 2:
		plt.annotate(xlabels[plot_idx-1], xy=(40, 0.02), size=12)
	if plot_idx == 3:
		plt.annotate(xlabels[plot_idx-1], xy=(40, 0.04), size=12)
	if plot_idx == 3:
		plt.legend(loc='best', fontsize=9, frameon=False)
	if plot_idx == 1 or plot_idx == 2:
		plt.tick_params(labelbottom='off')
	if plot_idx == 2:
		plt.ylabel('Distribution (arb.)')
	if plot_idx == 3:
		plt.xlabel('Angle (deg)')

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