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

from pathlib import Path

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

try:
    from scm.input_classes import AMS
except ImportError:
    from scm.input_classes.drivers import AMS


WORKDIR = "01-run_workdir"


def adf_settings(task: str, xc_type: str, xc_name: str) -> Settings:
    settings = Settings()
    settings.input.AMS.Task = task
    if task == "SinglePoint":
        settings.input.AMS.Properties.OrbitalsInfo = "Yes"
    settings.input.ADF.Basis.Type = "DZP"
    settings.input.ADF.NumericalQuality = "Basic"
    setattr(settings.input.ADF.XC, xc_type, xc_name)
    AMS.from_settings(settings)
    return settings


def run_job(name: str, system: ChemicalSystem, settings: Settings) -> AMSJob:
    job = AMSJob(name=name, molecule=system, settings=settings)
    print(f"\n--- {name} input ---")
    print(job.get_input())
    print(f"--- end {name} input ---\n")
    result = job.run()
    if not result.ok():
        raise RuntimeError(f"AMS job failed: {name}")
    return job


def main() -> None:
    init(folder=WORKDIR)
    try:
        outdir = Path("figures")
        outdir.mkdir(exist_ok=True)

        ethylene = ChemicalSystem.from_smiles("C=C")
        opt_settings = adf_settings("GeometryOptimization", "GGA", "PBE")
        opt_job = run_job("ethylene_opt_pbe_dzp", ethylene, opt_settings)

        optimized = opt_job.results.get_main_system()
        view(
            optimized,
            guess_bonds=True,
            direction="along_pca3",
            width=500,
            height=360,
            picture_path=str(outdir / "optimized_ethylene.png"),
        )

        gga_sp = run_job(
            "ethylene_sp_pbe_dzp",
            optimized,
            adf_settings("SinglePoint", "GGA", "PBE"),
        )
        hybrid_sp = run_job(
            "ethylene_sp_b3lyp_dzp",
            optimized,
            adf_settings("SinglePoint", "Hybrid", "B3LYP"),
        )

        print("Finished jobs:")
        for job in (opt_job, gga_sp, hybrid_sp):
            print(f"{job.name}: {job.path}")
    finally:
        finish()


if __name__ == "__main__":
    main()