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

import os
from dataclasses import dataclass

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


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


ISOMERS = (
    Isomer("butan-1-ol", "alcohol", "CCCCO", "butan_1_ol"),
    Isomer("butan-2-ol", "alcohol", "CCC(O)C", "butan_2_ol"),
    Isomer("2-methylpropan-1-ol", "alcohol", "CC(C)CO", "2_methylpropan_1_ol"),
    Isomer("2-methylpropan-2-ol", "alcohol", "CC(C)(C)O", "2_methylpropan_2_ol"),
    Isomer("ethoxyethane (diethyl ether)", "ether", "CCOCC", "ethoxyethane"),
    Isomer(
        "1-methoxypropane (methyl n-propyl ether)",
        "ether",
        "COCCC",
        "1_methoxypropane",
    ),
    Isomer(
        "2-methoxypropane (methyl isopropyl ether)",
        "ether",
        "COC(C)C",
        "2_methoxypropane",
    ),
)


def conformer_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.pre = "export OMP_NUM_THREADS=1"
    return settings


def validate_settings(settings: Settings) -> None:
    input_text = ConformersJob(settings=settings).get_input()
    InputParser().to_dict("conformers", input_text)


def main() -> None:
    for variable in ("AMSHOME", "AMSBIN"):
        if not os.environ.get(variable):
            raise RuntimeError(f"Required environment variable {variable} is not set")
    if not (os.environ.get("SCMLICENSE") or os.environ.get("SCM_CLOUD_CREDS")):
        raise RuntimeError("SCMLICENSE or SCM_CLOUD_CREDS must be set")

    settings = conformer_settings()
    validate_settings(settings)

    init(folder="01-conformers_workdir")
    try:
        for isomer in ISOMERS:
            system = ChemicalSystem.from_smiles(isomer.smiles)
            job = ConformersJob(
                name=isomer.slug,
                molecule=system,
                settings=settings.copy(),
            )
            job.run()
            if not job.ok():
                raise RuntimeError(f"Conformer calculation failed for {isomer.name}")
            n_unique = len(job.results.get_conformers())
            print(f"{isomer.name}: {n_unique} unique optimized conformer(s)")
    finally:
        finish()


if __name__ == "__main__":
    main()