#!/usr/bin/env python

#this script set up and runs several molecular dynamics trajectories
#this script must be called in the directory containing the n traj folders named as 000n
#a properly formatted Job file is required in the same foulder of the script

#################################################
from os import path, getcwd 			#
from os import chdir, system, mkdir		#
from time import sleep				#
from shutil import copy, rmtree, copytree	#
import numpy as np				#
#################################################

# Define calculation parameters #################
TimeLimit               ='100'			# in hours: "TimeLimit:00:00"
Ntraj			=10000			#
Njobtraj		=5000			# Make sure Ntraj / Njobtraj is an integer

#Ei / eV	E0 / eV	Alpha / m/s	Tn / K	Theta / degrees
beam_parameters = np.array(
[[0.94,	0.93,	158.,	296., 0.],	# Shirhatti 2016
[1.18,	1.14,	245.,	400., 0.],
[1.29,	1.24,	207.,	500., 0.],
[1.55,	1.49,	292.,	620., 0.],
[1.8,	1.73,	323.,	740., 0.],
[2.12,	2.03,	364.,	910., 0.],
[2.56,	2.47,	371.,	1060, 0.],
[0.275,	0.277,	52.,	296., 0.],	# Rahinov 2008
[0.317,	0.318,	60.,	296., 0.],	# Nozzle temperature is assumed to 296 K (not sure what it was)
[0.447,	0.449,	67.,	296., 0.],	# E0 is transformed from v0 using 0.5mv^2
[0.52,	0.524,	48.,	296., 0.],
[0.779,	0.779,	114.,	296., 0.],
[0.974,	0.979,	98.,	296., 0.],
[1.268,	1.278,	81.,	296., 0.],
[0.49,	1.18,	273,	1060., 50.],	# Shirhatti 2016, with different theta incidence angles
[0.69,	1.18,	273,    1060., 40.],
[1.03,	2.43,	321,    1060., 50.],
[1.27,	2.38,	322,    1060., 44.],
[1.47,	2.43,	321,    1060., 40.],
[1.76,	2.38,	322,    1060., 32.]]
)
#Beam_idx = [7, 8, 9, 10, 11, 12, 13]	# indices of the beam parameters to be used
#Beam_idx = [0, 1, 2, 3, 4, 5, 6]
Beam_idx = [14, 15, 16, 17, 18, 19] # different theta angles
#Beam_idx.append( 6 )
#Beam_idx = [1]
Surface = 'ms'				# Surface is frozen (fs), distorted (ds), or mobile (ms)
State_selected = 1			# =0 v, j, mj selected; =1 v, j; =2 v; =3 rovibrational states according to beam temperature
RoVib = [0, 2, 0]			# nu, j and mj state selected, according to State_selected
parallel_momentum = True		# Whether to include parallel momentum when theta != 0, or not
beam_parameters[:,3] = 1060.	#DONT FORGET TO REMOVE

#################################################
if (Ntraj % Njobtraj) != 0.:
	print('WARNING: Ntraj/Njobtraj is not an integer')
	exit()
Njobs = int(Ntraj/Njobtraj)

if State_selected > 3 or State_selected < 0:
	print('WARNING: State selection is wrong')
	exit()

if State_selected == 0 and abs(RoVib[2]) > RoVib[1]:
	print('WARNING: mJ is too large')
	exit()

here=getcwd()				#workdir
for j in Beam_idx:
	if State_selected == 0:
		if beam_parameters[j][4] != 0.:
			if parallel_momentum == False:
				newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}mj{5:d}_T{6:3.1f}_nopxy'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1], RoVib[2], beam_parameters[j][4] )
			else:
				newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}mj{5:d}_T{6:3.1f}'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1], RoVib[2], beam_parameters[j][4] )
		else:
			newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}mj{5:d}'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1], RoVib[2] )
	elif State_selected == 1:
		if beam_parameters[j][4] != 0.:
			if parallel_momentum == False:
				newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}_T{5:3.1f}_nopxy'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1], beam_parameters[j][4] )
			else:
				newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}_T{5:3.1f}'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1], beam_parameters[j][4] )
		else:
			newdir = '{0:s}/{1:3.2f}_{2:s}_v{3:d}j{4:d}'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1] )
	elif State_selected == 2:
		if beam_parameters[j][4] != 0.:
			if parallel_momentum == False:
				newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}_v{4:d}_T{5:3.1f}_nopxy'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3], RoVib[0], beam_parameters[j][4] )
			else:
				newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}_v{4:d}_T{5:3.1f}'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3], RoVib[0], beam_parameters[j][4] )
		else:
			newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}_v{4:d}'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3], RoVib[0] )
	else:
		if beam_parameters[j][4] != 0.:
			if parallel_momentum == False:
				newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}_T{4:3.1f}_nopxy'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3], beam_parameters[j][4] )
			else:
				newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}_T{4:3.1f}'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3], beam_parameters[j][4] )
		else:
			newdir = '{0:s}/{1:3.2f}_{2:s}_TN{3:4.0f}'.format( here, beam_parameters[j][0], Surface, beam_parameters[j][3] )

	if path.exists(newdir):
		rmtree( newdir )
	copytree( 'standard_{0:s}'.format( Surface ), newdir )

	chdir( newdir + '/DynamicsINPUTS' )

	input_f = open('in.inp', 'w')
	input_f.write('{0:d}         #n_trac number of trajectories\n'.format( Ntraj ))
	input_f.write('14.172944914  #zstart  z_start in a.u.\n')
	input_f.write('{0:d}             #state selected (=0 v, j, mj selected; =1 v, j; =2 v; =3 rovibrational states according to beam temperature)\n'.format( State_selected ))
	input_f.write('{0:d}             #qvib    vibrational state\n'.format( RoVib[0] ))
	input_f.write('{0:d}             #ql      rotational quantum number l\n'.format( RoVib[1] ))
	input_f.write('{0:d}             #qm      rotational quantum number l\n'.format( RoVib[2] ))
	input_f.write('{0:f}          #BEAM ENERGY in eV\n'.format( beam_parameters[j][0] ))
	input_f.write('gasphase.dat  #potfile asymtotic potential file\n')
	input_f.write('{0:f}          #BEAM TEMPERATURE\n'.format( beam_parameters[j][3] ))
	input_f.write('1             #TON SWITCH (TON=0 No TRANSLATIONAL ENERGY Distribution | TON=1 TRANSLATIONAL ENERGY DISTRIBUTION)\n')
	input_f.write('{0:f}         #alpha in m/s\n'.format( beam_parameters[j][2] ))
	input_f.write('{0:f}          #Estream in eV\n'.format( beam_parameters[j][1] ))
        input_f.write('{0:f}            #Theta angle of incidence in degrees\n'.format( beam_parameters[j][4] ))
	if parallel_momentum:
		input_f.write('true          #Add parallel momentum when theta != 0\n')
	else:
		input_f.write('false          #Add parallel momentum when theta != 0\n')
	input_f.write('431001        #seed number for random number (directory)\n')
	input_f.write('10234789      #seed for random number 2      (slab number)\n')
	input_f.close()

	system('./gen_poscar in.inp output')
#	system('./gen_poscar_noparallelmotion in.inp output')
	chdir( newdir )
	system('./create_poscar.py')

	for i in range(Njobs):
		Nfirst = '{0:06d}'.format( i*Njobtraj + 1)
		Nlast = '{0:06d}'.format( int(Nfirst) + Njobtraj - 1 )

		modify = 'sed -e \"s|_FIRSTJOB_|%s|g\" Job_mob_ > Job' % Nfirst 	# modify the Jobfile
		system(modify)								#
		modify = 'sed -i \"s|_TIMELIMIT_|%s|g\" Job' % TimeLimit                #
		system(modify)
		modify = 'sed -i \"s|_LASTJOB_|%s|g\" Job' % Nlast
		system(modify)
		if State_selected:
			modify = 'sed -i \"s|_NAME_|{0:3.2f}{1:s}{2:d}{3:d}|g\" Job'.format( beam_parameters[j][0], Surface, RoVib[0], RoVib[1] )
		else:
			modify = 'sed -i \"s|_NAME_|{0:3.2f}{1:s}TN|g\" Job'.format( beam_parameters[j][0], Surface )
		system(modify)

		system('qsub Job')	# RUN JOB #######################################
		sleep(5)

	chdir( here )