import os

# This PLAMS script locates TS for several reactions using the Nudged Elastic Band (NEB) method. 
# For each reaction two xyz files are required: one corresponding to the reactant state, called 
# '{}_initial.xyz', and one corresponding to the product state, called '{}_final.xyz' (e.g. 
# 'MyMolecule_initial.xyz' and 'MyMolecule_final.xyz')

# The folder containing the xyz files:
xyz_folder = os.path.join(os.environ['AMSHOME'],'examples','AMS','NEB_PLAMS','xyz')

names = [name.rsplit('_initial.xyz')[0] for name in os.listdir(xyz_folder) if name.endswith('_initial.xyz')]

# Settings for the AMS driver
neb_sett = Settings()
neb_sett.input.ams.Task = 'NEB'
neb_sett.input.ams.Properties.NormalModes = 'Yes'
neb_sett.input.ams.GeometryOptimization.Convergence.Step = 1.0e-3

# Settings for the engine (here we use the DFTB engine with semiempirical GFN1-xTB method)
engine_sett = Settings()
engine_sett.input.DFTB.Model='GFN1-xTB'

for name in sorted(names):

	# For NEB we need two system blocks in the AMS input (the initial and final molecule). 
	# In PLAMS you can have multiple system blocks by passing the the AMSJob a dictionary of molecules.
	# The 'keys' of the dictionary will be used as the headers of the System block
	mols = {}
	# The initial molecule should be in the main 'System' block (the main system block has no header). 
	# The key of the mols dictionary should therefore be an empty string
	mols[''] = Molecule(os.path.join(xyz_folder,name+'_initial.xyz'))
	# The final molecule should be specified in a system block with the header 'final'
	mols['final'] = Molecule(os.path.join(xyz_folder,name+'_final.xyz'))

	# Create and run the job:
	job = AMSJob(molecule=mols, name=name, settings=neb_sett+engine_sett)
	job.run()
	
	print('')
	print("System name: {}".format(name))
	if job.ok():
		pes_point_character = job.results.readrkf('AMSResults', 'PESPointCharacter', file='NEB_TS_final')
		print("NEB calculation converged. PES Point character: {}".format(pes_point_character))
		print("Left  TS barrier: {0:.6f} [Hartree]".format(job.results.readrkf('NEB', 'LeftBarrier')))
		print("Right TS barrier: {0:.6f} [Hartree]".format(job.results.readrkf('NEB', 'RightBarrier')))
	else:
		print("Unsuccesful NEB calculation")
	print('')