#!/usr/bin/env amspython
from scm.plams import *
from scm.params import *
import numpy as np
import os


def main():
    engine_collection = get_engine_collection()
    get_jobs_and_training_set(engine_collection, engine_id="UFF")
    get_parameter_interface()


def get_engine_collection():
    ec = EngineCollection()

    s = Settings()
    s.input.ReaxFF.ForceField = "CHON-2019.ff"
    s.input.ReaxFF.TaperBO = "Yes"
    ec.add_entry("ReaxFF:CHON-2019", Engine(s))

    s = Settings()
    s.input.ForceField.Type = "UFF"
    ec.add_entry("UFF", Engine(s))

    s = Settings()
    s.input.MLPotential.Model = "ANI-2x"
    ec.add_entry("ANI-2x", Engine(s))

    s = Settings()
    s.input.DFTB.Model = "GFN1-xTB"
    ec.add_entry("DFTB", Engine(s))

    s = Settings()
    s.input.ADF.XC.GGA = "PBE"
    s.input.ADF.Basis.Type = "TZP"
    s.input.ADF.Unrestricted = "No"
    ec.add_entry("PBE_TZP", Engine(s))

    s = Settings()
    s.input.ADF.XC.GGA = "PBE"
    s.input.ADF.Basis.Type = "TZP"
    s.input.ADF.Unrestricted = "Yes"
    s.input.ADF.SpinPolarization = 1.0
    ec.add_entry("PBE_TZP_spin1", Engine(s))

    return ec


def get_jobs_and_training_set(engine_collection: EngineCollection, engine_id: str = "UFF"):
    # initialize ResultsImporter
    ri = ResultsImporter(settings={"units": {"energy": "kcal/mol"}})
    ri.job_collection.engines = engine_collection

    # directory with xyz files e_3.xyz, p_3.xyz, e_9.xyz, p_9.xyz etc.
    xyzdir = os.path.expandvars("$AMSHOME/scripting/scm/params/examples/benchmark/ISOL6")

    # id vs. reference-value (kcal/mol)
    reactions = {"3": 9.77, "9": 21.76, "10": 6.83, "13": 33.52, "14": 5.3}

    for rid, reference in reactions.items():
        reactant = Molecule(f"{xyzdir}/e_{rid}.xyz")
        product = Molecule(f"{xyzdir}/p_{rid}.xyz")
        # add the jobs to the job collection
        # and reaction energy to training set
        ri.add_reaction_energy(reactants=[reactant], products=[product], normalization="r0", reference=reference)

        # access the added jobs and set the corresponding engine id
        for jid in ri.data_set[-1].jobids:
            ri.job_collection[jid].extra_engine = engine_id

    ri.store(backup=False, store_settings=False)


def get_parameter_interface():
    interf = NoParameters()
    interf.yaml_store("parameter_interface.yaml")


if __name__ == "__main__":
    main()