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

from dataclasses import dataclass
from pathlib import Path

import scm.plams as plams
from scm.base import ChemicalSystem
from scm.conformers import ConformersJob
from scm.input_classes import Conformers
from scm.plams import Settings


@dataclass(frozen=True)
class MoleculeSpec:
    name: str
    smiles: str
    slug: str


MOLECULES: tuple[MoleculeSpec, ...] = (
    MoleculeSpec("butan-1-ol", "CCCCO", "butan-1-ol"),
    MoleculeSpec("butan-2-ol", "CCC(O)C", "butan-2-ol"),
    MoleculeSpec("2-methylpropan-1-ol", "CC(C)CO", "2-methylpropan-1-ol"),
    MoleculeSpec("2-methylpropan-2-ol", "CC(C)(C)O", "2-methylpropan-2-ol"),
    MoleculeSpec("ethoxyethane (diethyl ether)", "CCOCC", "ethoxyethane"),
    MoleculeSpec("1-methoxypropane (methyl n-propyl ether)", "COCCC", "1-methoxypropane"),
    MoleculeSpec("2-methoxypropane (methyl isopropyl ether)", "COC(C)C", "2-methoxypropane"),
)


def build_settings() -> Settings:
    settings = Settings()
    settings.input.ams.Task = "Generate"
    settings.input.ams.Generator.Method = "RDKit"
    settings.input.ams.Generator.RDKit.InitialNConformers = 8
    settings.input.MLPotential.Model = "AIMNet2-wB97MD3"
    settings.runscript.nproc = 1
    settings.runscript.preamble_lines = ["export OMP_NUM_THREADS=1"]
    Conformers.from_settings(settings)
    return settings


def run_job(spec: MoleculeSpec, settings: Settings) -> ConformersJob:
    molecule = ChemicalSystem.from_smiles(spec.smiles)
    job = ConformersJob(name=spec.slug, molecule=molecule, settings=settings.copy())
    job.run()
    if not job.ok():
        raise RuntimeError(f"Conformers job failed for {spec.name}")
    energies = job.results.get_relative_energies("kcal/mol")
    print(f"{spec.name}: {len(energies)} conformers kept; lowest relative energy = {energies[0]:.6f} kcal/mol")
    return job


def main() -> None:
    plams.init(folder="01-run_workdir")
    settings = build_settings()
    for spec in MOLECULES:
        run_job(spec, settings)
    plams.finish()


if __name__ == "__main__":
    main()