#!/usr/bin/env python
import sys
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 = [ '0.94', '1.18', '1.29', '1.55', '1.80', '2.12', '2.56' ]

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)):
		traj= '%.6d' % 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:
			angles.append( [float(post_dat[25].split()[0]), float(post_dat[27].split()[0]),	float(post_dat[25].split()[2]), float(post_dat[27].split()[2])] )
		elif 'SCATTERING' in outcome:
			angles_scat.append( [float(post_dat[24].split()[0]), float(post_dat[26].split()[0])] )
		elif 'TRAPPING' in outcome:
			angles_trap.append( [float(post_dat[24].split()[0]), float(post_dat[26].split()[0])] )

	chdir(here)

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

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(4):
		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(2):
		mu.append( np.mean(angles_scat[:,i]) )
		sigma.append( np.std(angles_scat[:,i], ddof=1) )
		error.append( sigma[i+4] / len(angles_scat)**0.5 )
	for i in range(2):
		mu.append( np.mean(angles_trap[:,i]) )
		sigma.append( np.std(angles_trap[:,i], ddof=1) )
		error.append( sigma[i+4] / len(angles_trap)**0.5 )
	
	print( 'Reacted (t=0):' )
	print( "theta\t\tphi" )
	print( '{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})'.format(mu[0], error[0], sigma[0], mu[1], error[1], sigma[1]) )

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

	print( 'Scattered (t=0):' )
	print( "theta\t\tphi" )
	print( '{:.1f} +/- {:.1f} ({:.1f})\t{:.1f} +/- {:.1f} ({:.1f})'.format(mu[4], error[4], sigma[4], mu[5], error[5], sigma[5]) )

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

if False:
	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=2
Ncols=1
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(3.37,4.))
#mpl.rc("font", size=10)
mpl.rc("text", usetex=True)

xlabels = [r'$\theta$', r'$\phi$']

for plot_idx in range(1,Nrows*Ncols + 1):
	plt.subplot(Nrows,Ncols,plot_idx)

	if len(angles) != 0:
		n, bins = np.histogram( angles[:,plot_idx-1], bins=np.arange(0.,183.*plot_idx,3.*plot_idx), 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 )
		if plot_idx == 1:
			av = 0.
		else:
			av = (n[0] + n[-1]) / 2.
		n = np.insert(n, 0, av)
		n = np.append(n, av)

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

		n, bins = np.histogram( angles[:,plot_idx+1], bins=np.arange(0.,183.*plot_idx,3.*plot_idx), 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 )
		if plot_idx == 1:
			av = 0.
		else:
			av = (n[0] + n[-1]) / 2.
		n = np.insert(n, 0, av)
		n = np.append(n, av)

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

	if len(angles_scat) != 0:
		angles_scat = np.array(angles_scat)
		n, bins = np.histogram( angles_scat[:,plot_idx-1], bins=np.arange(0.,183.*plot_idx,3.*plot_idx), 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 )
		if plot_idx == 1:
			av = 0.
		else:
			av = (n[0] + n[-1]) / 2.
		n = np.insert(n, 0, av)
		n = np.append(n, av)

		xnew = np.linspace(0.,180.*plot_idx,180.*plot_idx)
		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)
		n, bins = np.histogram( angles_trap[:,plot_idx-1], bins=np.arange(0.,183.*plot_idx,3.*plot_idx), density=True )
#		n, bins = np.histogram( angles_trap[:,plot_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 )
		if plot_idx == 1:
			av = 0.
		else:
			av = (n[0] + n[-1]) / 2.
		n = np.insert(n, 0, av)
		n = np.append(n, av)

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

	if plot_idx == 1:
		x = np.linspace(0.,180.,180.)
		n = np.sin( np.radians(x) )
		n /= sum(n)
		plt.plot(x,n, label='Statistical', linewidth=1, color='k')
	else:
		x = np.linspace(0.,360.,360.)
		n = np.array( 360*[1.] )
		n /= sum(n)
		plt.plot(x,n, label='Statistical', linewidth=1, color='k')

	y_limits = plt.ylim()
	if plot_idx == 1:
		plt.plot( [115., 115.], plt.ylim(), linestyle='dotted', c='k' )
	if plot_idx == 2:
		plt.plot( [1., 1.], [0., y_limits[1]], linestyle='dotted', c='k' )
		plt.plot( [121., 121.], [0., y_limits[1]], linestyle='dotted', c='k' )
		plt.plot( [241., 241.], [0., y_limits[1]], linestyle='dotted', c='k' )
	plt.ylim(0., y_limits[1])

	plt.tick_params(labelleft='off')
	plt.xlim(0,180*plot_idx)
	if plot_idx == 2:
		plt.xticks( np.arange(0., 420., 60.) )
	if plot_idx == 1:
		plt.annotate(xlabels[plot_idx-1], xy=(160, y_limits[1]*0.8), size=14)
	if plot_idx == 2:
		plt.annotate(xlabels[plot_idx-1], xy=(320, y_limits[1]*0.5), size=14)
	if plot_idx == 1:
		plt.legend(loc='best', fontsize=8, frameon=False)
	plt.ylabel('Distribution (arb.)')
	plt.xlabel('Angle (deg)')

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