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

mpl.rc("font", size=20)

first = 1
last  = 500

MC           = 12.000
MO           = 16.000
MH           = 1.008 
Dt           = 0.4E-15
Masses       = [ MC ] + [ MO ] + 4 * [ MH ]

Etransfer = []
here=getcwd()				#workdir
for i in range(first,(last+1)):
	traj= '%.4d' % i		# define working directory
	newdir=path.join(here, traj)	# updating mypath

	if not path.exists(newdir): 
		print "folder", newdir, "not found. exiting..."
		quit()

	chdir(newdir)

	if path.exists('outcome'):
		outcome = open('outcome', 'r').readlines()
		if not 'SCATTERING\n' in outcome: continue
		Etotal = []
		for POSCARName in ["initial", "final"]:
			lines = open( 'EnergyTransfer/POSCAR_' + str(POSCARName), "r" ).readlines()

			Velocities = []

			for i in range( 16, 22 ):
				Velocity = [ float( lines[i].split()[0] ), float( lines[i].split()[1] ), float( lines[i].split()[2] ) ]
				Velocities.append( Velocity )

			DOFs = 18
			Ekin = 0.

			for i in range( 6 ):
				for j in range( 3 ):
				        Ekin = Ekin + 0.5 * Velocities[ i][ j ] ** 2. * Masses[ i ] * 1822.88839 / 0.52917721092**2. / 41.341373337**2. * 27.21138505

			lines = open( 'EnergyTransfer/OSZICAR_' + str(POSCARName), 'r' ).readlines()
			Efree = float( lines[-1].split()[2] )
			Etotal.append(Efree + Ekin)

		Etransfer.append( Etotal[0] - Etotal[1] )

n, bins, patches = plt.hist(Etransfer, bins=30, facecolor='b', normed='True')
#bins = np.linspace(-3.5,0., num=100)

mu, sigma = norm.fit(Etransfer)
y = mpl.mlab.normpdf( bins, mu, sigma)
l = plt.plot(bins, y, 'r--', linewidth=4, color='blue', label='Cu(111)')

print mu
print sigma

#plt.xlim(-0.5,1.25)
#plt.ylim(0,0.05)
#plt.title('Energy transfer of CHD3 v1=1 161 kJ/mol')
plt.xlabel('Initial - Final (eV)')
plt.ylabel('Relative distribution (arbitrary units)')
plt.legend(loc='upper left')

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