#!/usr/bin/env python3

#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               ='72'			# in hours: "TimeLimit:00:00"
Ntraj			=2000			#
Njobtraj		=50			# Make sure Ntraj / Njobtraj is an integer
Ncores			=1			#

#Ei / eV	v0 / m/s	Alpha / m/s	Tn / K
beam_parameters = np.array(
[[0.004, 626.5, 55.9, 293],	# Elham Pt(211) 2019 paper, table 4
[0.009, 943.5, 127.8, 293],	# for H2
[0.013, 1085.1, 111.6,293],
[0.014, 1145.2, 118.7,293],
[0.025, 1531.4, 96.6,293],
[0.035, 1447.5, 293.9,293],
[0.043, 2031.2, 80.6,293],
[0.132, 3392.1, 578.0,500],
[0.181, 3959.8, 690.8,700],
[0.169, 4009.0, 185.2,1300],
[0.233, 4442.8, 862.5, 900],
[0.282, 4838.8, 1022.9, 1100],
[0.338, 5223.2, 1215.6, 1300],
[0.413, 5617.0, 1535.8, 1500],
[0.454, 5790.7, 1711.3, 1700],
[0.0073, 573.8, 73.9, 300],    # for D2, from Charlotte
[0.0084, 615.9, 84.2, 300],
[0.0121, 738.3, 99.3, 300],
[0.0212, 980.4, 117.5, 300],
[0.0304, 1176.7, 143.5, 300],
[0.0437, 1423.7, 139.4, 300],
[0.0530, 1568.1, 151.4, 300],
[0.0692, 1795.3, 163.4, 300],
[0.0889, 2022.9, 220.6, 400],
[0.1078, 2212.3, 278.4, 500],
[0.1249, 2367.1, 330.8, 600],
[0.1405, 2521.1, 328.7, 700],
[0.1436, 2512.5, 405.5, 750],
[0.008, 626.8,  49.8,   293],   # for D2
[0.027, 1103.2, 134.9,  293],
[0.04,  1379.2, 75.5,   293],
[0.054, 1555.6, 218.1,  600],
[0.076, 1860.,  237.5,  800],
[0.11,  2239.9, 267.3,  1100],
[0.13,  2430.7, 311.9,  500],
[0.14,  2531.1, 294.5,  1100],
[0.234, 3191.,  548.1,  900],
[0.276, 3628.1, 160.3,  1300],
[0.346, 3814.8, 772.,   1300],
[0.457, 4304.1, 989.4,  1700]]
)
#Beam_idx = [4, 5, 6, 7, 8, 10, 11, 12, 13, 14] # indices of the beam parameters to be used
#Beam_idx = [14, 13, 12, 11, 10, 8, 7, 6, 5, 4]
#Beam_idx = [15,16,17,18,19,20,21,22,23,24,25,26,27]
Beam_idx = [15,19,27]
Theta = 2.8				# Incidence angle (degrees)
Phi = 0.                                # Orientation of velocity vector in XY plane (degrees)
Surface = 'fs'				# Surface is frozen (fs), distorted (ds), or mobile (ms)
State_selected = False			# Vibrational state population according to nozzle temperature (False), or state selected (True) where mj is sampled randomly
RoVib = [1, 1]				# nu and j state selected, if State_selected == True
Nlayers = 5
Nsupercellx = 1
Nsupercelly = 3
mil_h = 17
mil_k = 15
mil_l = 15

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

here=getcwd()				#workdir
for j in Beam_idx:
	if State_selected:
		newdir = '{0:s}/{1:4.3f}_{2:s}_v{3:d}j{4:d}'.format( here, beam_parameters[j][0], Surface, RoVib[0], RoVib[1] )
	else:
		newdir = '{0:s}/{1:4.3f}_{2:s}_TN'.format( here, beam_parameters[j][0], Surface )

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

	chdir( newdir + '/DynamicsINPUTS/inicond_D2_FWHM' )

	input_f = open('Input/input.txt', 'w')
	input_f.write('2.014            # mass1 (atomic mass units)\n')
	input_f.write('2.014            # mass2 (atomic mass units)\n')
	input_f.write('{0:f}               # nozzle temperature (K)\n'.format( beam_parameters[j][3] ))
	input_f.write('1000               # number of grids for single-body Hamiltonian\n')
	input_f.write('6                  # maximum number of vibrational states "Nvs"\n')
	input_f.write('50                 # maximum angular momentum quantum number "Jmax"\n')
	input_f.write('1.0                # threshold energy for the harmonic fit (eV)\n')
	input_f.write('{0:f}            # average initial translational energy (eV)\n'.format( beam_parameters[j][0] ))
	input_f.write('{0:d}              # number of classical trajectories\n'.format( Ntraj ))
	input_f.write('0.046196059      # broadening of normal translational energy distribution\n')
	input_f.write('1                  # ortho-para states:  1="Yes", 0="No"\n')
	input_f.write('0.666666666666667  # ortho-para ratio for even J\n')
	input_f.write('0.333333333333333  # ortho-para ratio for odd J\n')
	input_f.write('1                  # velocity distribution in the molecular beam:  0:"E_k=<E>", 1:"f(E_k)=max(f)"\n')
	if State_selected:
		State_int = 1
	else:
		State_int = 0
	input_f.write('{0:d}                  # state-resolved calculations 1="Yes", 0="No"\n'.format( State_int ))
	input_f.write('{0:d}                  # v (if state-resolved flag is on)\n'.format( RoVib[0] ))
	input_f.write('{0:d}                  # j (if state-resolved flag is on)\n'.format( RoVib[1] ))
	input_f.write('1                  # Are alpha and v0 specified? 1="Yes", 0="No"\n')
	input_f.write('{0:f}                # alpha (m/s)\n'.format( beam_parameters[j][2] ))
	input_f.write('{0:f}               # stream velocity (m/s)\n'.format( beam_parameters[j][1] ))
	input_f.write('0.8                # rotational temperature (as a fraction of Tn)\n')
	input_f.write('{0:f}                 # incidence angle Theta (degrees)\n'.format( Theta ))
	input_f.write('{0:f}                 # incidence angle Phi (degrees)\n'.format( Phi ))
	input_f.close()

	system('./geninicon_1.py')
	chdir( newdir )
	if Surface == 'fs':
		system('./create_poscar.py {:} {:} {:} {:} {:} {:} {:}'.format( Nlayers, Nsupercellx, Nsupercelly, mil_h, mil_k, mil_l, 'False' ))
	else:
		system('./create_poscar.py {:} {:} {:} {:} {:} {:} {:}'.format( Nlayers, Nsupercellx, Nsupercelly, mil_h, mil_k, mil_l, 'True' ))

	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_ > 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)
		modify = 'sed -i \"s|_NCORES_|%s|g\" Job' % Ncores
		system(modify)
		if State_selected:
			modify = 'sed -i \"s|_NAME_|{0:4.3f}{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:4.3f}{1:s}TN|g\" Job'.format( beam_parameters[j][0], Surface )
		system(modify)

		system('sbatch Job')	# RUN JOB #######################################
#		sleep(2)

	chdir( here )