#!/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 import interpolate

#folders = [ '2.56_ms_v2j2_angles', '2.56_ms_v2j8_angles' ]

Nangle = 100
Ns = 40
Minr = 1.2
Maxr = 2.2
here=getcwd()
def collectangles(folder):
	angles = np.zeros((Nangle,Ns))
	chdir(folder)

	Ntraj = int( popen('ls | egrep "^[0-9].....$" | wc -l').read() )
	Nreac = 0
	Ntot = 0
	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').readlines()
		if not 'UNCLEAR' in outcome[0]:
			Ntot += 1
		if 'REACTION ' in outcome[0]:
			Nreac += 1
			data = open('analysis_rcom.dat', 'r').readlines()
			for j in range( 1, len(data) ):
				r = float( data[j].split()[1] )
				z = float( data[j].split()[2] )
				if z > 2.5: continue
				if r < Minr or r > Maxr: continue
				idx0 = round( float( data[j].split()[8] ) / ( 180. / float(Nangle-1) ) )
				idx1 = round( (r-Minr) / ( (Maxr-Minr) / float(Ns-1) ) )
				angles[idx0][idx1] += 1

	S0 = float(Nreac) / float(Ntot)
	for j in range( Ns ):
		angles[:,j] /= sum( angles[:,j] )
		angles[:,j] *= S0

	chdir(here)
	
	return angles

Nrows=4
Ncols=2
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(3.69,5.))
#mpl.rc("font", size=10)
mpl.rc("text", usetex=True)

angles_v0j2 = collectangles('2.56_ms_v0j2_angles')
angles_v0j8 = collectangles('2.56_ms_v0j8_angles')
angles_v2j2 = collectangles('2.56_ms_v2j2_angles')
angles_v2j8 = collectangles('2.56_ms_v2j8_angles')

X = np.linspace(0., 180., Nangle)

ax2 = plt.subplot(Nrows,Ncols,1)

idx = 4
plt.plot( X, angles_v0j2[:,idx], color='r', label=r'$\nu=0, J=2$' )
plt.plot( X, angles_v0j8[:,idx], color='b', label=r'$\nu=0, J=8$' )

ylim = plt.ylim()
plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(a)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

ax2.set_title(r'$\nu=0$', loc='center')

plt.tick_params(labelbottom=False, labelleft=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.yticks([0.,0.01])
plt.xticks(np.linspace(0.,150.,6))
#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.xlabel('Reaction path')
#plt.ylabel('Distribution (arb.)')

ax2 = plt.subplot(Nrows,Ncols,2)

idx = 4
plt.plot( X, angles_v2j2[:,idx], color='r', label=r'$\nu=2, J=2$' )
plt.plot( X, angles_v2j8[:,idx], color='b', label=r'$\nu=2, J=8$' )

ylim = plt.ylim()
#plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(b)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

ax2.set_title(r'$\nu=2$', loc='center')

plt.tick_params(labelbottom=False, labelleft=False, labelright=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.xticks(np.linspace(0.,180.,7))
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.ylabel('Distribution (arb.)')
#plt.xlabel(r'$\theta$ angle (degrees)')

ax2 = plt.subplot(Nrows,Ncols,3)

idx = 10
plt.plot( X, angles_v0j2[:,idx], color='r', label=r'$\nu=0, J=2$' )
plt.plot( X, angles_v0j8[:,idx], color='b', label=r'$\nu=0, J=8$' )

ylim = plt.ylim()
plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(c)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=False, labelleft=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.xticks(np.linspace(0.,150.,6))
#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.xlabel('Reaction path')
plt.ylabel('Sticking probability')

ax2.yaxis.set_label_coords(-0.25,0.)

plt.subplot(Nrows,Ncols,4)

idx = 10
plt.plot( X, angles_v2j2[:,idx], color='r', label=r'$\nu=2, J=2$' )
plt.plot( X, angles_v2j8[:,idx], color='b', label=r'$\nu=2, J=8$' )

ylim = plt.ylim()
#plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(d)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=False, labelleft=False, labelright=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.yticks([0.,0.01])
plt.xticks(np.linspace(0.,180.,7))
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.ylabel('Distribution (arb.)')
#plt.xlabel(r'$\theta$ angle (degrees)')

plt.subplot(Nrows,Ncols,5)

idx = 20
plt.plot( X, angles_v0j2[:,idx], color='r', label=r'$\nu=0, J=2$' )
plt.plot( X, angles_v0j8[:,idx], color='b', label=r'$\nu=0, J=8$' )

ylim = plt.ylim()
plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(e)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=False, labelleft=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.xticks(np.linspace(0.,150.,6))
#plt.legend(loc='best', numpoints=1, frameon=True)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.xlabel(r'$\theta$ angle (degrees)')
#plt.ylabel('Distribution (arb.)')

plt.subplot(Nrows,Ncols,6)

idx = 20
plt.plot( X, angles_v2j2[:,idx], color='r', label=r'$\nu=2, J=2$' )
plt.plot( X, angles_v2j8[:,idx], color='b', label=r'$\nu=2, J=8$' )

ylim = plt.ylim()
#plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(f)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=False, labelleft=False, labelright=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.yticks([0., 0.01])
plt.xticks(np.linspace(0.,180.,7))
#plt.legend(loc='lower left', numpoints=1, frameon=True, fontsize=6)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.ylabel('Distribution (arb.)')
#plt.xlabel(r'$\theta$ angle (degrees)')

plt.subplot(Nrows,Ncols,7)

idx = 35
plt.plot( X, angles_v0j2[:,idx], color='r', label=r'$\nu=0, J=2$' )
plt.plot( X, angles_v0j8[:,idx], color='b', label=r'$\nu=0, J=8$' )

ylim = plt.ylim()
plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(g)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=True, labelleft=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.xticks(np.linspace(0.,150.,6))
#plt.legend(loc='best', numpoints=1, frameon=True)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
plt.xlabel(r'$\theta$ angle (degrees)')
#plt.ylabel('Distribution (arb.)')


plt.subplot(Nrows,Ncols,8)

idx = 35
plt.plot( X, angles_v2j2[:,idx], color='r', label=r'$J=2$' )
plt.plot( X, angles_v2j8[:,idx], color='b', label=r'$J=8$' )

ylim = plt.ylim()
#plt.annotate('$r={:3.2f}$'.format( Minr + (Maxr-Minr)/(Ns-1)*idx ) + r' \r{A}', xy=(5., 0.85*ylim[1]), size=10)
plt.annotate('(h)', xy=(155., 0.85*ylim[1]), size=10)

plt.plot([117.,117.], [0., ylim[1]], color='k', linestyle='--')

plt.tick_params(labelbottom=True, labelleft=False, labelright=True, direction='in', top=True, right=True)
#plt.yticks([])
plt.yticks([0.,0.01,0.02,0.03])
plt.xticks(np.linspace(0.,180.,7))
plt.legend(loc='lower left', numpoints=1, frameon=True, fontsize=6)
plt.xlim(0.,180.)
plt.ylim( 0., ylim[1] )
#plt.ylabel('Distribution (arb.)')
plt.xlabel(r'$\theta$ angle (degrees)')

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