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

from scm.plams import AMSJob, Settings, finish, from_smiles, init, log, packmol

from common import completed_jobs, reaxff_settings


def main() -> None:
    reusable = completed_jobs("01-initial-system_workdir", "setup_opt")
    if reusable:
        job = reusable["setup_opt"]
        log(f"Reusing completed ReaxFF geometry optimization at {job.path}.")
        return

    log("Building neutral glycine NH2-CH2-COOH from SMILES NCC(=O)O.")
    glycine = from_smiles("NCC(=O)O")
    water = from_smiles("O")
    log(f"Glycine has {len(glycine)} atoms and formal charge {glycine.properties.get('charge', 0)}.")
    log("Packing one glycine and 40 water molecules at an overall density of 1.0 g/cm^3.")
    packed = packmol(
        molecules=[glycine, water],
        n_molecules=[1, 40],
        density=1.0,
        region_names=["glycine", "water"],
        keep_bonds=True,
        seed=38040,
    )
    log(f"Packed system has {len(packed)} atoms, regions glycine and water, "
        f"and density {packed.get_density() / 1000:.6f} g/cm^3.")

    settings = reaxff_settings()
    settings.input.ams.Task = "GeometryOptimization"
    settings.input.ams.GeometryOptimization.Convergence.Quality = "Basic"
    job = AMSJob(name="setup_opt", molecule=packed, settings=settings)
    log("Starting ReaxFF bad-contact removal with GeometryOptimization Convergence Quality Basic.")
    job.run()
    if not job.ok():
        raise RuntimeError(f"Initial ReaxFF geometry optimization failed: {job.results.get_errormsg()}")
    log(f"Initial ReaxFF geometry optimization succeeded at {job.path}.")


if __name__ == "__main__":
    init(folder="01-initial-system_workdir")
    try:
        main()
    finally:
        finish()