#!/usr/bin/env amspython
import os
import time

from ase.io import read
from scm.plams import *
from scm.plams.interfaces.adfsuite.ase_calculator import AMSCalculator
from sella import Sella

"""

This examples runs a transition state search with the Sella python package and
compares results and timings to the builtin transition state search in AMS (using a DFTB hessian to initialize the DFT transition state search).

This example requires an ADF license.

"""

def adf_settings():
    s = Settings()
    s.input.adf.basis.type = 'DZP'
    s.input.adf.basis.core = 'None'
    s.input.adf.xc.gga = 'pbe'
    s.input.adf.symmetry = 'nosym'
    return s

def run_sella(molecule):
    """ Run a TS search with Sella """
    s = adf_settings()
    s.input.ams.task = 'SinglePoint'
    s.input.ams.properties.gradients = 'Yes'

    atoms = toASE(molecule)

    if os.path.exists("sella.log"):
        os.remove("sella.log")
    if os.path.exists("sella.traj"):
        os.remove("sella.traj")

    with AMSCalculator(settings=s, amsworker=True) as calc:
        atoms.calc = calc
        opt = Sella(atoms, internal=True, logfile="sella.log", trajectory="sella.traj")
        opt.run(fmax=0.027, steps=50)  # 0.027 eV/ang roughly matches AMS default criterion 0.001 Ha/ang

    traj = read("sella.traj", ":")
    print(f"Sella converged after {len(traj)} single-point calculations")
    print(f"Final energy: {atoms.get_potential_energy()} eV")

    optimized_mol = fromASE(atoms)
    return optimized_mol

def run_ams(molecule):
    """ Run DFT transition state search but use initial hessian calculated with DFTB """
    s = adf_settings()
    s.input.ams.task = 'TransitionStateSearch'
    s.input.ams.geometryoptimization.initialhessian.type = 'CalculateWithFastEngine'
    job = AMSJob(settings=s, molecule=molecule, name='ams_ts')
    job.run()
    print(f"AMS finished after {len(job.results.get_history_property('Energy'))} iterations")
    print(f"AMS optimized energy: {job.results.get_energy(unit='eV')} eV")

def get_molecule():
    return AMSJob.from_input("""
        System
          atoms
             C         0.049484    0.042994    0.000000
             H        -0.068980    0.638928   -0.915972
             H        -0.068980    0.638928    0.915972
             H        -0.841513   -0.626342    0.000000
             H         0.555494   -1.148227    0.000000
             Hg        2.303289   -0.007233    0.000000
             Cl        4.429752    0.776056    0.000000
             Cl        1.342057   -2.676083    0.000000
          end
        End

    """).molecule['']

def main():
    os.environ['OMP_NUM_THREADS'] = '1'
    init()
    mol = get_molecule()

    start = time.time()
    run_sella(mol)
    print(f"Sella finished in {time.time()-start} seconds")

    start = time.time()
    run_ams(mol)
    print(f"AMS finished in {time.time()-start} seconds")

    finish()

if __name__ == '__main__':
    main()