#!/usr/bin/env amspython
from scm.plams import *
from scm.params import *

def main():
    """
    run_reference : Runs a volume scan on LiF, calculating the stress tensor at each point.

    add_engines: Sets up per-job engines defining different k-space settings

    add_volume_scan : Adds the volume scan points to the job collection and training set, and links the job collection entries to the previously defined engines

    add_enthalpy_of_formation: Adds jobs and training set entries for the experimental LiF enthalpy of formation

    add_bond_length: Adds a training set entry for the optimized bond length of the F2 molecule
    """
    # If you have already run the reference job, simply set pesscan_rkf to the path to the ams.rkf file
    pesscan_rkf = run_reference('LiF_primcell.xyz')
    #pesscan_rkf = './plams_workdir/volumescan.results/ams.rkf'

    results_importer = ResultsImporter()
    add_engines(results_importer) # in order to control k-space sampling on a per-job basis (needed for xTB on periodic systems)
    add_volume_scan(results_importer, pesscan_rkf)
    add_enthalpy_of_formation(results_importer)
    add_bond_length(results_importer)
    results_importer.store(backup=False, store_settings=False)

    generate_parameter_interface()


def run_reference(xyz_file):
    """ 
        Sets up a volume scan of LiF with PBEsol-D3(BJ).
        For each point, the stress tensor is calculated.
        Returns the path to the ams.rkf file

        NOTE: This is an expensive calculation that might take several hours.
    """
    init()
    mol = Molecule(xyz_file)
    s = Settings()
    s.input.band.KSpace.Type = 'Symmetric'
    s.input.band.KSpace.Quality = 'Normal' # N.B. this should be set to 'Good' for production calculations!
    s.input.band.XC.GGA = 'PBEsol'
    s.input.band.XC.Dispersion = 'Grimme3 BJDamp'
    s.input.band.Basis.Type = 'DZP'
    s.input.ams.Task = 'pesscan'
    s.input.ams.PESScan.ScanCoordinate.NPoints = 5
    s.input.ams.PESScan.ScanCoordinate.CellVolumeScalingRange = '0.80 1.20'
    s.input.ams.PESScan.CalcPropertiesAtPESPoints = 'Yes' # needed to get the stress tensor at each point
    s.input.ams.GeometryOptimization.MaxIterations = 0
    s.input.ams.GeometryOptimization.PretendConverged = 'Yes'
    s.input.ams.Properties.StressTensor = 'Yes'
    job = AMSJob(settings=s, molecule=mol, name='volumescan')
    job.run()
    finish()

    return job.results.rkfpath()

def generate_parameter_interface():
    """ 
        Generate parameter_interface.yaml activating 
        the Li and F repulsive potentials for optimization 
    """
    # s is the *default* engine settings used during the parametrization
    # It is normal to leave it as an empty Settings and set the engine settings per job
    s = Settings() 
    interface = GFN1xTBParameters(settings=s)

    # Repulsive potential of Li.
    interface['Li:alpha'].is_active = True
    interface['Li:Z'].is_active = True

    # Repulsive potential of F.
    interface['F:alpha'].is_active = True
    interface['F:Z'].is_active = True

    interface.yaml_store("parameter_interface.yaml")

def add_engines(results_importer : ResultsImporter):
    """ Control the k-space settings for each job by adding different engine definitions. """
    s = Settings()
    s.input.dftb.KSpace.Type = 'Symmetric'
    s.input.dftb.KSpace.Quality = 'Normal'
    results_importer.job_collection.engines.add_entry('kspace_normal_symmetric', Engine(s))

    s = Settings()
    s.input.dftb.KSpace.Type = 'Symmetric'
    s.input.dftb.KSpace.Quality = 'Good'
    results_importer.job_collection.engines.add_entry('kspace_good_symmetric', Engine(s))

def add_volume_scan(results_importer : ResultsImporter, ref_rkf : str):
    """
        The volume scan is added using the add_pesscan_singlepoints() method, which supports extracting the stress tensor.

        If no stress tensor is desired, an alternative is to instead use the add_singlejob() method with task='PESScan'

        See the documentation for more details.
    """
    results_importer.add_pesscan_singlepoints(ref_rkf, properties={
        'relative_energies': {
            'relative_to': 'min',
            'unit': 'eV',
        },
        'stresstensor_diagonal_3d': {
            'unit': 'GPa',
        },
    }, extra_engine='kspace_normal_symmetric')


def add_enthalpy_of_formation(results_importer : ResultsImporter) :
    """
        Adds experimental enthalpy of formation of LiF
        The jobs are added with Good k-space quality, since total energies of different
        supercells are sensitive to the k-space sampling. (Moreover Li is a metal)
    """
    LiF = Molecule('LiF_primcell.xyz')
    Li = Molecule('Li_primcell.xyz')
    F2 = Molecule('F2.xyz')
    results_importer.add_reaction_energy(
        reactants=[Li, F2],
        reactants_names=['Li_bulk', 'F2'],
        products=[LiF],
        products_names=['LiF_bulk'],
        normalization='p0', # reaction energy per LiF (first product)
        task='GeometryOptimization',
        reference=-6.394,
        unit='eV',
        extra_engine='kspace_good_symmetric' # applied to the Li_bulk and LiF_bulk jobs. The F2 molecule is nonperiodic so the k-space settings are ignored
    )

def add_bond_length(results_importer : ResultsImporter):
    """ add experimental F2 bond length. This requires that the job F2 has been defined (which is done in add_enthalpy_formation) """
    results_importer.data_set.add_entry('distance("F2",0,1)', reference=1.41)


if __name__ == '__main__':
    main()