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

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

	data = open('DynamicsINPUTS/initial_pos_momenta.info', 'r').readlines()
	data2 = open('DynamicsINPUTS/init_vals.dat', 'r').readlines()
	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

		#if int( data2[(i-1)*4].split()[5] ) != 0: continue
		outcome = open('outcome', 'r').readlines()
		if not 'UNCLEAR' in outcome[0]:
			theta = np.degrees( float( data[i-1].split()[4] ) )
			ptheta = float( data[i-1].split()[10] ) * 0.529177249 / 1822.888486209 * 21.876913 # angstrom to bohr, electron mass to amu, velocity in Hartree atomic units to angstrom / fs
			weight = 1. / float( np.sin( np.radians( theta ) ) )

			angles_stat.append( np.array([theta, ptheta, weight]) )
			if 'REACTION ' in outcome[0]:
				angles.append( np.array([theta, ptheta, weight]) )

	angles_stat = np.array( angles_stat )
	#angles_stat_binned, edges_stat = np.histogramdd( angles_stat, density=False, bins=4 )
	angles_stat_binned, xedges, yedges = np.histogram2d( angles_stat[:,1], angles_stat[:,0], range=[[min(angles_stat[:,1]), max(angles_stat[:,1])], [20., 160.]], bins=36, density=True, weights=angles_stat[:,2] )
	# for i in range(edges[1].size - 1):
		# if np.sum(angles_stat_binned[:,i]) == 0: continue
		# angles_stat_binned[:,i] /= np.sum(angles_stat_binned[:,i])

	angles = np.array( angles )
	#angles_binned, edges = np.histogramdd( angles, density=False, bins=edges_stat )
	#angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[-99., 99.], [-10., 170.]], bins=[xedges, yedges], density=False, weights=angles[:,2] )
	angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], range=[[min(angles[:,1]), max(angles[:,1])], [20., 160.]], bins=[xedges, yedges], density=True, weights=angles[:,2] )
	#angles_binned, xedges, yedges = np.histogram2d( angles[:,1], angles[:,0], bins=[xedges, yedges], density=False )
	#for i in range(yedges.size - 1):
		#if np.sum(angles_binned[i,:]) == 0: continue
		#angles_binned[i,:] /= np.sum(angles_binned[i,:])

	angles_binned -= angles_stat_binned
	angles_binned += 1.

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

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

def plotangles(idx, title, folder):
	ax2 = plt.subplot(Nrows,Ncols,idx)

	angles, edges = collectangles(folder)

	vlimit = max( abs( np.amin(angles) ), abs( np.amax(angles) ) )
	plt.pcolormesh(edges[1], edges[0], angles, vmin=1.-(vlimit-1.), vmax=vlimit, cmap='bwr')
	#cbar = plt.colorbar(label='Relative intensity', ticks=[1.])
	#cbar.ax.set_yticklabels(['Avg.'])
	cbar = plt.colorbar(label='Relative intensity', ticks=[1.-(vlimit-1.),1.,vlimit])
	cbar.ax.set_yticklabels(['Low', 'Avg.', 'High'])

	y = ( plt.ylim()[1] - plt.ylim()[0] ) * 0.05 + plt.ylim()[0]
	plt.annotate(title, xy=(0.525, 0.075), xycoords='axes fraction', size=10, color='k', bbox=dict(boxstyle='round', fc='w'))

	plt.xlabel(r'$\theta$ angle (degrees)')
	plt.ylabel(r'$p_\theta$ (amu\r{A}$^2$/fs)')

	#plt.xticks([0., 45., 90., 135., 180.])
	plt.xticks([20., 55., 90., 125., 160.])

	return

nu = 0
for i in range(1,9):
	plotangles( i, r'$\nu={0:d},J={1:d}$'.format(nu, i), '2.56_ms_v{0:d}j{1:d}'.format(nu, i) )

plt.tight_layout()
plt.savefig('momentum_theta_relative_v{:d}.pdf'.format(nu))