#!/usr/bin/env python3

#################################################
from os import path
import numpy as np
#from numpy import array, dot, degrees, clip, arccos	#
#from numpy.linalg import norm, inv              #
#################################################

#####################################################################################################
### P A R A M E T E R S #############################################################################
#####################################################################################################

Write_analysis = False	# Whether to write the analysis files or not
R_CRIT_min	= 2.0
T_CRIT_min	= 100.0	# If molecule is larger than R_CRIT_min but smaller than R_CRIT_max for at least T_CRIT_min fs, the molecule is considered to be reacted
R_CRIT_max	= 2.2
R_WARNING	= 3.5
Z_CRIT		= 6.0

tstep	=   0.1	# i-pi stride=10# [ fs ] 
T_MAX	=  1000	# [ fs ]
	
MC		=  12.000 # [ amu ]
MH		=   1.008 # [ amu ]
MD		=   2.014 # [ amu ]
Masses	= [ MC, MH, MH, MH, MH ]

Mass	= sum(Masses)	# methane mass
surface_atoms = 45
molecule_atoms = 5

#########################################
					#
au2Ang  =    0.52917721092		# conversion constants
amu2au  = 1822.88839			#
fs2au   =   41.341373337		#
au2eV   =   27.21138505			#
deg2rad =    0.0174532925		#
J2eV    =    6.24181      * 10**18	#
h       =    6.6260695729 * 10**-34	#
					#
Kb      = 8.6173857e-05                 # [ eV/K ] Boltzmann  
amu2kg  = 1.6605402e-27			#
ang2m   = 1e-10				#
fs2s    = 1e-15				#
eV2J    = 1.60217733e-19		#
					#
#########################################

def com( atoms_list, masses_list ):

        CoM = [ 0., 0., 0. ]
        if len(atoms_list) != len(masses_list):
                print( "number of atoms and masses must be the same" )
                quit()

        for i in range(len(atoms_list)):
                for j in range(3):
                        CoM[j] += masses_list[i] * atoms_list[i][j]

        for j in range(3): CoM[j] = CoM[j]/sum(masses_list)

        return CoM

def ab(v1, v2):
        " Returns the angle in radians between vectors 'v1' and 'v2'  "

        def unit_vector(vector):
                "Returns the unit vector of the vector."
                return vector / np.linalg.norm(vector)

        v1_u = unit_vector(v1)
        v2_u = unit_vector(v2)

        return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))

def distance(A,B):
        "return the distance between two 3D points"
        return np.linalg.norm(np.asarray(A)-np.asarray(B))

def kinetic(v, M):
        v_tot = np.linalg.norm( v )
        Ki = 0.5 * M * ( v_tot**2 ) * amu2kg * ang2m**2 / fs2s**2 / eV2J 
        return Ki

#####################################################################################################

# This part is slightly hardcoded with some assumptions
lines = open('simulation.pos_centroid.xyz', 'r')
CellLines = lines.readline()		# read file line by line
CellLines = lines.readline()		# read file line by line
a = [ float(CellLines.split()[2]), 0., 0. ]
b = [ float(CellLines.split()[3]) * np.cos( np.radians( float( CellLines.split()[7] ) ) ), float(CellLines.split()[3]) * np.sin( np.radians( float( CellLines.split()[7] ) ) ), 0. ]
c = [ 0., 0., float(CellLines.split()[4]) ]
Cell = np.array([a, b, c])			# transformation matrix
Cell_= np.linalg.inv(Cell)

### Positions ########################################################################################

C  = [ ]                # position of the atoms 
H1 = [ ]                # for each timestep
H2 = [ ]                #
H3 = [ ]                #
H4 = [ ]                #

#-----------------------------------------------------------------------------------------------------
def ExtractPositions(posfile):
	header = 2
	block = surface_atoms + molecule_atoms + header	# surf atoms + mol atoms + headline

	lines = open(posfile).readlines()
	steps = int( len( lines ) / block )
	for i in range(steps):
		C.append ( np.array([ float(x) for x in lines[( header + i*block + surface_atoms     )].split()[1:4] ]) )
		H1.append( np.array([ float(x) for x in lines[( header + i*block + surface_atoms + 1 )].split()[1:4] ]) )
		H2.append( np.array([ float(x) for x in lines[( header + i*block + surface_atoms + 2 )].split()[1:4] ]) )
		H3.append( np.array([ float(x) for x in lines[( header + i*block + surface_atoms + 3 )].split()[1:4] ]) )
		H4.append( np.array([ float(x) for x in lines[( header + i*block + surface_atoms + 4 )].split()[1:4] ]) )
#-----------------------------------------------------------------------------------------------------

ExtractPositions('simulation.pos_centroid.xyz')
C = np.array(C)
H1 = np.array(H1)
H2 = np.array(H2)
H3 = np.array(H3)
H4 = np.array(H4)
steps = len(C)

### Real Traj #################################################################################
### Find the real trajectory of the atmos considering the eventual supercell 
### supercell border crossing
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def Real_Trajectory(atom, ref):
	'''
	Atom is an array containing, for each step, [x, y, z] of the atom in direct coordinates.
	Ref is the reference atom in the molecule. Atom coordinates will be changed so that the real position of the 
	atom will be considered and not necessarily the periodic image in the [0,0,0] cell. 
	The reference atoms is needed to find out which cell to start from
	'''

	r_min	= 99999											# absolute distance between the reference and the atom
	diff	= np.zeros(3)									# component-wise distance 
	Real	= []

	for i in range(-1, 2):									# find out in which cell is the atom
		for j in range(-1, 2):								# in respect of the reference one
			dist = ref[0] - (atom[0] + np.array([i,j,0]))	#
			dist = np.linalg.norm(dist)						# this is the distance between the reference atom
															# and the periodic image in cell [i,j,0] of the tested atom
			if dist < r_min:								# store the periodic image closest to the reference atom
				r_min = dist								# at the beginning of the trajectory
				a_min = i
				b_min = j

	cell = np.array([a_min , b_min, 0])						# the cell considered is the real [0, 0, 0] + cell[a_min, b_min, 0]

	Real.append( atom[0] + cell )						 	# initial step

	for i in range(1, len(ref)):							# assume that the next step will
		Real.append(atom[i]+cell)							# be in the same cell
		for j in range(3):									# check if that's actually the case
			diff[j] = Real[i][j] - Real[(i-1)][j]
			if   diff[j] < -0.5:							# if it's not update the position
				cell[j]    += 1.							# and the reference cell accordingly
				Real[i][j] += 1.
			elif diff[j] >  0.5:
				cell[j]    -= 1.
				Real[i][j] -= 1.
	
	return np.array(Real)
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

### Cartesian to Direct ############################################################################

C_d = np.dot(C, Cell_)
H1_d = np.dot(H1, Cell_)
H2_d = np.dot(H2, Cell_)
H3_d = np.dot(H3, Cell_)
H4_d = np.dot(H4, Cell_)

### Get rid of periodicity ############################################################################

C_d = Real_Trajectory( C_d,  C_d )
H1_d = Real_Trajectory( H1_d,  C_d )
H2_d = Real_Trajectory( H2_d,  C_d )
H3_d = Real_Trajectory( H3_d,  C_d )
H4_d = Real_Trajectory( H4_d,  C_d )

### Direct to Cartesian ############################################################################

C = np.dot(C_d, Cell)
H1 = np.dot(H1_d, Cell)
H2 = np.dot(H2_d, Cell)
H3 = np.dot(H3_d, Cell)
H4 = np.dot(H4_d, Cell)

### Velocities ####################################################################################
#-----------------------------------------------------------------------------------------------------
def ExtractVelocities(velfile):
	'''
	Note that the velocities in the input file are given in m/s, whereas internally we work with angstrom/fs,
	hence the 1e-5 multiplication.
	'''
	header = 2
	block = surface_atoms + molecule_atoms + header	# surf atoms + mol atoms + headline

	lines = open(velfile).readlines()
	steps = int( len( lines ) / block )
	for i in range(steps):
		vC.append ( np.array([ float(x)*1e-5 for x in lines[( header + i*block + surface_atoms     )].split()[1:4] ]) )
		vH1.append( np.array([ float(x)*1e-5 for x in lines[( header + i*block + surface_atoms + 1 )].split()[1:4] ]) )
		vH2.append( np.array([ float(x)*1e-5 for x in lines[( header + i*block + surface_atoms + 2 )].split()[1:4] ]) )
		vH3.append( np.array([ float(x)*1e-5 for x in lines[( header + i*block + surface_atoms + 3 )].split()[1:4] ]) )
		vH4.append( np.array([ float(x)*1e-5 for x in lines[( header + i*block + surface_atoms + 4 )].split()[1:4] ]) )
#-----------------------------------------------------------------------------------------------------
vC = []
vH1 = []
vH2 = []
vH3 = []
vH4 = []

ExtractVelocities('simulation.vel_centroid.xyz')
vC = np.array(vC)
vH1 = np.array(vH1)
vH2 = np.array(vH2)
vH3 = np.array(vH3)
vH4 = np.array(vH4)

### COM  ########################################################################################

COM  = []	# position of the CoM
vCOM = []	# velocity of the CoM
KCOM = []	# kinetic E of the CoM 

for i in range(steps):
	CenterOfMass_x = Masses[0]*C[i] + Masses[1]*H1[i] + Masses[2]*H2[i] + Masses[3]*H3[i] + Masses[4]*H4[i]
	CenterOfMass_v = Masses[0]*vC[i] + Masses[1]*vH1[i] + Masses[2]*vH2[i] + Masses[3]*vH3[i] + Masses[4]*vH4[i]
	
	COM.append( CenterOfMass_x/Mass )
	vCOM.append( CenterOfMass_v/Mass )

	k_ics		= kinetic(vCOM[-1][0], Mass)
	k_ypsilon	= kinetic(vCOM[-1][1], Mass)
	k_zed		= kinetic(vCOM[-1][2], Mass)
	k_total		= kinetic(vCOM[-1][:], Mass)

	KCOM.append( [ k_ics , k_ypsilon , k_zed, k_total ] )
KCOM = np.array(KCOM)

### Bond Lengths ##################################################################################
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def rMol(A1, A2):
	printWARNING = False
	r = []
	for k in range( len(A1) ):
		r.append( distance(A1[k], A2[k]) )
		if r[-1] > R_WARNING:
			open('BONDveryLONG', 'w').write("\n")
	r = np.array(r)
	
	return r
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

rCH1 = rMol(C, H1)
rCH2 = rMol(C, H2)
rCH3 = rMol(C, H3)
rCH4 = rMol(C, H4)

### Evaluate Outcome #################################################################################
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def reacted(radius, label, R_CRIT_min, R_CRIT_max, T_CRIT_min):
	'''
	Given a vector containing the radius (bond distance) at each timestep, it returns True if the bond has been broken.
	the bond is considered broken if it is larger than R_CRIT_max or if it stays larger than R_CRIT_min 
	for T_CRIT_min timesteps
	'''

	BACK     = False					# True if the molecule back reacts: r > R_CRIT_min for t < T_CRIT_min
	REACT    = False
	THRESHOLD= np.ceil( T_CRIT_min / tstep )
	steps    = len(radius)
	d_steps  = 0
	react_t  = 999999.

	for k in range(steps):
		if d_steps > 0 and radius[k] < R_CRIT_min:
			BACK    = True
			d_steps = 0 

		elif radius[k] > R_CRIT_max:	
			REACT   = True 

		elif radius[k] > R_CRIT_min and d_steps < THRESHOLD:
			if d_steps == 0: react_t = k * tstep
			d_steps += 1

		elif d_steps >= THRESHOLD:
			REACT = True

	return label, REACT, BACK, react_t

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

reaction = []
reaction.append( reacted( rCH1, 'H1', R_CRIT_min, R_CRIT_max, T_CRIT_min ) )
reaction.append( reacted( rCH2, 'H2', R_CRIT_min, R_CRIT_max, T_CRIT_min ) )
reaction.append( reacted( rCH3, 'H3', R_CRIT_min, R_CRIT_max, T_CRIT_min ) )
reaction.append( reacted( rCH4, 'H4', R_CRIT_min, R_CRIT_max, T_CRIT_min ) )

SCATTERING = False 
BACK_REACTION = False
REACTION = False
TRAPPING = False

N_bond_broken = 0
for i in range(4):
	if reaction[i][1] == True:
		N_bond_broken += 1
		bond_broken	= reaction[i][0]
		react_t		= reaction[i][3]
		REACTION	= True

	if reaction[i][2] == True:
		BACK_REACTION = True		

if REACTION == False:
	for i in range(steps):
		if COM[i][2] > Z_CRIT and vCOM[i][2] > 0:
			SCATTERING = True
			scat_t = i*tstep
			break

if REACTION == False and SCATTERING == False:
	if steps*tstep > T_MAX:
		TRAPPING = True

### Analysis #########################################################################################

### Angles ###########################################################################################
def Angles( DISS, H1, H2, H3, C ):
	'''
	Returns theta, beta and gamma.
	'''
	THETA = []
	BETA = []
	GAMMA = []

	for j in range( len(DISS) ):
		CH_bond 	= DISS[j] - C[j]					# dissociating bond
		baricentro 	= np.array( com( [ H1[j], H2[j], H3[j] ], [1., 1., 1.] ) )	# umbrella geometric center == com if the masses are all the same
		umbrella   	= C[j] - baricentro

		THETA.append( np.degrees( ab( CH_bond, [0,0,1] ) ) )
		BETA.append( np.degrees( ab( umbrella, [0,0,1] ) ) )
		GAMMA.append( np.degrees( ab( CH_bond , umbrella ) ) )

	return THETA, BETA, GAMMA

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

if Write_analysis == True:

	if REACTION and bond_broken == 'H1':
		theta, beta, gamma = Angles(H1, H2, H3, H4, C)
	elif REACTION and bond_broken == 'H2':
		theta, beta, gamma = Angles(H2, H1, H3, H4, C)
	elif REACTION and bond_broken == 'H3':
		theta, beta, gamma = Angles(H3, H2, H1, H4, C)
	elif REACTION and bond_broken == 'H4':
		theta, beta, gamma = Angles(H4, H2, H3, H1, C)
	else:
		theta, beta, gamma = Angles(H1, H2, H3, H4, C)

	output_angles = open( 'angles.dat', 'w' )
	output_angles.write('t        theta  beta    gamma\n')
	for i in range( steps ):
		output_angles.write( '{:7.2f}  {:4.1f}  {:4.1f}  {:4.1f}\n'.format( (i*tstep), theta[i], beta[i], gamma[i] ) )
	output_angles.close()

	output_r = open('analysis_r.dat', 'w')
	output_c = open('analysis_com.dat', 'w')
	output_r.write('t        r_CH1  r_CH2  r_CH3  r_CH4\n')
	output_c.write('t        Z       X        Y        vZ         Kz\n')

	for i in range(steps):
		output_r.write('{:7.2f}  {:4.3f}  {:4.3f}  {:4.3f}  {:4.3f}\n'.format( (i*tstep), rCH1[i], rCH2[i], rCH3[i], rCH4[i] ))
		output_c.write('{:7.2f}  {:6.4f}  {:6.4f}  {:6.4f}  {:8.6f}  {:8.6f}\n'.format( (i*tstep), COM[i][2], COM[i][0], COM[i][1], vCOM[i][2], KCOM[i][2] ))
	output_r.close()
	output_c.close()

### Reaction outcome ################################################################################

if [ REACTION, TRAPPING, SCATTERING ].count(True) > 1:
	open('outcome', 'w').write('WARNING: MULTIPLE OUTCOMES DETECTED\n')
	open('WARNING', 'w').write('WARNING: MULTIPLE OUTCOMES DETECTED\n')
	open('WARNING', 'a').write('REACT, BACK, TRAP, SCAT == {:s}, {:s}, {:s}, {:s}\n'.format( REACTION, BACK_REACTION, TRAPPING, SCATTERING ))

elif N_bond_broken > 1:
	open('outcome', 'w').write('WARNING: %d critical bond lengths\n' % N_bond_broken )
	open('WARNING', 'w').write('WARNING: %d critical bond lengths\n' % N_bond_broken )

elif REACTION and N_bond_broken == 1:
	if BACK_REACTION:
		open('outcome', 'w').write('BACK-REACTION - {:s}\nafter {:6.2f} fs\n'.format( bond_broken, react_t ))
	else:
		open('outcome', 'w').write('REACTION - {:s}\nafter {:6.2f} fs\n'.format( bond_broken, react_t ))

elif TRAPPING:
	open('outcome', 'w').write('TRAPPING\ntotal propagation time: {:6.2f}\n'.format( steps*tstep ) )

elif SCATTERING:
	if BACK_REACTION:
		open('outcome', 'w').write('BACK-SCATTERING\nafter {:6.2f} fs\n'.format( scat_t ))
	else:
		open('outcome', 'w').write('SCATTERING\nafter {:6.2f} fs\n'.format( scat_t )) 

else:
	open('outcome', 'w').write('UNCLEAR \n')