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

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

def reactioncoordinate( r, Z, idx ):
	if len(r) != len(Z):
		print('lengths of r and Z are not equal')
		exit()
	
	s = [ 0. ]
	for i in range( 1, len( r ) ):
		dZ = Z[i] - Z[i-1]
		dr = r[i] - r[i-1]
		s.append( s[-1] + np.sqrt( dZ**2 + dr**2 ) )
	ds = s[idx]
	for i in range( len(s) ):
		s[i] -= ds
	
	return s

Nangle = 100
Ns = 40
s_limit = [-8.,0.]
here=getcwd()
def collectangles(folder):
	angles = []
	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').readlines()
		if 'REACTION ' in outcome[0]:
			t0 = round( float( outcome[1].split()[1] ) / 0.4 )
			
			data = open('analysis_rcom.dat', 'r').readlines()
			
			r = []
			Z = []
			theta = []
			#for j in range(1, t0+1):
			for j in range(1, len(data)):
				r.append( float( data[j].split()[1] ) )
				Z.append( float( data[j].split()[2] ) )
				theta.append( float( data[j].split()[8] ) )
			s = reactioncoordinate( r, Z, t0 )
			for j in range( len(s) ):
				#if s[j] < s_limit[0] or s[j] > s_limit[1]: continue

				angles.append( [s[j], theta[j]] )
	
	angles = np.array(angles)
	#angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[0., 180.], [s_limit[0], s_limit[1]]], bins=(Ns, Nangle), density=True )
	angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[0., 180.], [s_limit[0], s_limit[1]]], bins=(Nangle, Ns), density=True )
	
	#for j in range( Ns ):
	#	angles_binned[:,j] /= sum( angles_binned[:,j] )

	chdir(here)
	
	return angles_binned, [xedges, yedges]

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

vmin = 0.
vmax = 0.003

xticks = [-8., -4., 0.]

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

angles, edges = collectangles('2.56_ms_v0j2_angles')
plt.pcolormesh(edges[1], edges[0], angles, cmap='jet', vmin=vmin, vmax=vmax)

plt.plot(s_limit, [117.,117.], c='k', ls='--')
plt.annotate('$J=2$', xy=(0.65, 0.05), xycoords='axes fraction', size=10, color='white')

ax2.set_title(r'$\nu=0$', loc='center')
plt.yticks([0.,45.,90.,135.,180.])
plt.tick_params(labelbottom=False)
#plt.xlabel(r'Reaction path (\r{A})')
plt.ylabel(r'$\theta$ angle (degrees)')
ax2.yaxis.set_label_coords(-0.3,0.)

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

angles, edges = collectangles('2.56_ms_v2j2_angles')
plt.pcolormesh(edges[1], edges[0], angles, cmap='jet', vmin=vmin, vmax=vmax)

plt.plot(s_limit, [117.,117.], c='k', ls='--')
plt.annotate('$J=2$', xy=(0.65, 0.05), xycoords='axes fraction', size=10, color='white')

ax2.set_title(r'$\nu=2$', loc='center')
plt.tick_params(labelbottom=False, labelleft=False, left=False)
#plt.xlabel('Reaction path')
#plt.ylabel(r'$\theta$ angle (degrees)')

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

angles, edges = collectangles('2.56_ms_v0j8_angles')
plt.pcolormesh(edges[1], edges[0], angles, cmap='jet', vmin=vmin, vmax=vmax)

plt.plot(s_limit, [117.,117.], c='k', ls='--')
plt.annotate('$J=8$', xy=(0.65, 0.05), xycoords='axes fraction', size=10, color='white')

plt.yticks([0.,45.,90.,135.])
plt.xticks(xticks)
#plt.xlabel('Reaction path')
#plt.ylabel(r'$\theta$ angle (degrees)')

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

angles, edges = collectangles('2.56_ms_v2j8_angles')
cbar = plt.pcolormesh(edges[1], edges[0], angles, cmap='jet', vmin=vmin, vmax=vmax)

plt.plot(s_limit, [117.,117.], c='k', ls='--')
plt.annotate('$J=8$', xy=(0.65, 0.05), xycoords='axes fraction', size=10, color='white')

plt.xticks([-7.,-4.,0.])
plt.tick_params(labelleft=False, left=False)
plt.xlabel(r'Reaction coordinate (\r{A})')
ax2.xaxis.set_label_coords(0.,-0.28)
#plt.ylabel(r'$\theta$ angle (degrees)')

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)

cbar_ax = fig.add_axes([0.78, 0.2, 0.025, 0.685]) #xmin, ymin, xwidth, ywidth
fig.colorbar(cbar, cax=cbar_ax) #, ticks=np.arange(0.0,0.7,0.1)
cbar_ax.set_ylabel('Probability density')
fig.subplots_adjust(right=0.76, left=0.15)

plt.savefig('angles_reactionpath.pdf')
#plt.show()