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

from dataclasses import dataclass

from scm.base import ChemicalSystem
from scm.input_classes import AMS
from scm.plams import AMSJob, Settings, config, init


SYSTEM_BLOCK = """
System
   Atoms
      C 0 0 1.5
      C 0 0 4.5
      C 0 1.4202816622064793 1.5
      C 1.23 0.7101408311032397 4.5
   End
   Lattice
      2.46 0 0
      -1.23 2.130422493309719 0
      0 0 6
   End
End
"""


@dataclass(frozen=True)
class FunctionalSpec:
    key: str
    label: str
    input_dft: str
    pp_family: str
    pp_functional: str
    qe_vdw_corr: str | None = None
    d4_functional: str | None = None


FUNCTIONALS: tuple[FunctionalSpec, ...] = (
    FunctionalSpec(
        key="pbe_psl_paw",
        label="PBE / pslibrary-PAW",
        input_dft="pbe",
        pp_family="pslibrary-PAW",
        pp_functional="PBE",
    ),
    FunctionalSpec(
        key="b86bpbe_xdm_psl_paw",
        label="B86bPBE-XDM / pslibrary-PAW",
        input_dft="b86bpbe",
        pp_family="pslibrary-PAW",
        pp_functional="PBE",
        qe_vdw_corr="XDM",
    ),
    FunctionalSpec(
        key="r2scan_dojo",
        label="r2SCAN / Dojo",
        input_dft="r2scan",
        pp_family="Dojo",
        pp_functional="PBE",
    ),
    FunctionalSpec(
        key="r2scan_d4_dojo",
        label="r2SCAN-D4 / Dojo",
        input_dft="r2scan",
        pp_family="Dojo",
        pp_functional="PBE",
        d4_functional="R2SCAN",
    ),
)


SCAN_SEGMENTS: tuple[tuple[str, float, float, float], ...] = (
    ("seg1", 6.0, 9.0, 0.1),
    ("seg2", 9.0, 16.0, 0.5),
)


def npoints(start: float, stop: float, step: float) -> int:
    return int(round((stop - start) / step)) + 1


def build_settings(spec: FunctionalSpec, c_start: float, c_stop: float, step: float) -> Settings:
    settings = Settings()
    settings.input.ams.Task = "PESScan"
    settings.input.ams.PESScan.Optimize = "No"
    settings.input.ams.PESScan.CalcPropertiesAtPESPoints = "No"
    settings.input.ams.PESScan.ScanCoordinate.nPoints = npoints(c_start, c_stop, step)
    settings.input.ams.PESScan.ScanCoordinate.LatticeCRange = f"{c_start:.10g} {c_stop:.10g}"

    settings.input.QuantumESPRESSO.Pseudopotentials.Family = spec.pp_family
    settings.input.QuantumESPRESSO.Pseudopotentials.Functional = spec.pp_functional
    settings.input.QuantumESPRESSO.System.input_dft = spec.input_dft
    settings.input.QuantumESPRESSO.System.occupations = "Smearing"
    settings.input.QuantumESPRESSO.System.smearing = "Gaussian"
    settings.input.QuantumESPRESSO.System.degauss = 0.001
    settings.input.QuantumESPRESSO.System.ecutwfc = 100.0
    settings.input.QuantumESPRESSO.System.ecutrho = 1000.0
    if spec.qe_vdw_corr is not None:
        settings.input.QuantumESPRESSO.System.vdw_corr = spec.qe_vdw_corr

    settings.input.QuantumESPRESSO.K_Points._h = "automatic"
    settings.input.QuantumESPRESSO.K_Points._1 = "12 12 4 0 0 0"

    if spec.d4_functional is not None:
        settings.input.ams.EngineAddons.D4Dispersion.Enabled = "Yes"
        settings.input.ams.EngineAddons.D4Dispersion.Functional = spec.d4_functional

    AMS.from_settings(settings)
    return settings


def main() -> None:
    init(folder="01-run_workdir")
    config.log.stdout = 1

    system = ChemicalSystem(SYSTEM_BLOCK)

    for spec in FUNCTIONALS:
        for segment_name, c_start, c_stop, step in SCAN_SEGMENTS:
            job_name = f"{spec.key}_{segment_name}"
            settings = build_settings(spec, c_start, c_stop, step)
            job = AMSJob(molecule=system, settings=settings, name=job_name)
            print(f"Running {job_name}: {spec.label}, c = {c_start} -> {c_stop} A, step {step} A")
            result = job.run()
            if not result.ok():
                raise RuntimeError(f"Job failed: {job_name}")


if __name__ == "__main__":
    main()