#!/usr/bin/env amspython
from __future__ import annotations

from scm.base import ChemicalSystem, InputParser
from scm.plams import AMSJob, Settings, finish, init


def adf_settings(task: str, functional_kind: str, functional: str) -> Settings:
    settings = Settings()
    settings.input.ams.Task = task
    settings.input.adf.Basis.Type = "DZP"
    settings.input.adf.Basis.Core = "None"
    settings.input.adf.NumericalQuality = "Normal"
    settings.input.adf.SCF.Iterations = 100
    if functional_kind == "GGA":
        settings.input.adf.XC.GGA = functional
    elif functional_kind == "Hybrid":
        settings.input.adf.XC.Hybrid = functional
    else:
        raise ValueError(f"Unsupported functional kind: {functional_kind}")
    return settings


def validate_job(job: AMSJob) -> None:
    InputParser().to_dict("ams", job.get_input())


def main() -> None:
    init(folder="01-run_workdir")

    initial_system = ChemicalSystem.from_smiles("C=C")
    optimization = AMSJob(
        molecule=initial_system,
        settings=adf_settings("GeometryOptimization", "GGA", "PBE"),
        name="ethylene_opt_pbe_dzp",
    )
    validate_job(optimization)
    optimization.run()
    if not optimization.ok():
        raise RuntimeError("PBE/DZP geometry optimization failed")

    optimized_system = optimization.results.get_main_system()
    single_points = (
        ("ethylene_sp_pbe_dzp", "GGA", "PBE"),
        ("ethylene_sp_b3lyp_dzp", "Hybrid", "B3LYP"),
    )
    for name, functional_kind, functional in single_points:
        job = AMSJob(
            molecule=optimized_system,
            settings=adf_settings("SinglePoint", functional_kind, functional),
            name=name,
        )
        validate_job(job)
        job.run()
        if not job.ok():
            raise RuntimeError(f"Single-point calculation failed: {name}")

    finish()


if __name__ == "__main__":
    main()