import scm.plams

scm.plams.init(folder="NEB.results")

# ---------------------------------------------
# Engine Settings
# ---------------------------------------------
engine_sett = scm.plams.Settings()
engine_sett.input.MLPotential.Backend = "M3GNet"
engine_sett.input.MLPotential.Model = "M3GNet-UP-2022"
engine_sett.input.ams.EngineAddons.D3Dispersion.Enabled = "Yes"

# ---------------------------------------------
# Bulk Geometry Optimization
# ---------------------------------------------
mol = scm.plams.Molecule("LiTiS2_pristine.xyz")

sett = scm.plams.Settings()
sett.input.ams.Task = "GeometryOptimization"
sett.input.ams.GeometryOptimization.OptimizeLattice = "Yes"

job = scm.plams.AMSJob(settings=sett + engine_sett, molecule=mol, name="relax")
results = job.run()

if not job.ok():
    print("Bulk geometry optimization FAILED!")
    exit(0)

sett.input.ams.GeometryOptimization.OptimizeLattice = "No"

# ---------------------------------------------
# Initial State
# ---------------------------------------------
mol_relax = job.results.get_main_molecule()
coords_orig = mol_relax[25].coords

mol_init = mol_relax.copy()
mol_init.delete_atom(mol_init[25])

job = scm.plams.AMSJob(settings=sett + engine_sett, molecule=mol_init, name="init")
results = job.run()

if not job.ok():
    print("Initial state optimization FAILED!")
    exit(0)

mol_init_relax = job.results.get_main_molecule()

# ---------------------------------------------
# Final State
# ---------------------------------------------
mol_final = mol_init_relax.copy()
mol_final[1].coords = coords_orig

job = scm.plams.AMSJob(settings=sett + engine_sett, molecule=mol_final, name="final")
results = job.run()

if not job.ok():
    print("Final state optimization FAILED!")
    exit(0)

mol_final_relax = job.results.get_main_molecule()

# ---------------------------------------------
# NEB calculation
# ---------------------------------------------
sett.input.ams.Task = "NEB"
sett.input.ams.NEB.Images = 19
sett.input.ams.NEB.Iterations = 100

molecules = {"": mol_init_relax, "final": mol_final_relax}

job = scm.plams.AMSJob(settings=sett + engine_sett, molecule=molecules, name="NEB")
results = job.run()

if not job.ok():
    print("NEB calculation FAILED!")
    exit(0)

molecule_ts = results.get_main_molecule()

left_barrier = results.readrkf("NEB", "LeftBarrier")
right_barrier = results.readrkf("NEB", "RightBarrier")

left_barrier_eV = scm.plams.Units.convert(left_barrier, "au", "eV")
right_barrier_eV = scm.plams.Units.convert(right_barrier, "au", "eV")

print("Geometry of the TS:")
print(molecule_ts)
print(f"Left TS barrier : {left_barrier_eV:.6f} [eV]")
print(f"Right TS barrier: {right_barrier_eV:.6f} [eV]")

scm.plams.finish()