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

from pathlib import Path

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


SMILES = "c1ccc(cc1)C(=O)OOC(=O)c2ccccc2"
MODEL = "UMA-S-1.2-OMol"
TARGET_DISTANCE_ANGSTROM = 3.0
STEP_ANGSTROM = 0.1
WORKDIR = "01-benzoyl-peroxide-oo-scan_workdir"


def build_system() -> ChemicalSystem:
    system = ChemicalSystem.from_smiles(SMILES)
    system.guess_bonds()
    return system


def find_oo_bond(system: ChemicalSystem) -> tuple[int, int]:
    for atom_index, atom in enumerate(system.atoms):
        if atom.symbol != "O":
            continue
        for i, j, _bond in system.bonds.get_bonds_for_atom(atom_index):
            other_index = j if i == atom_index else i
            if system.atoms[other_index].symbol == "O":
                return tuple(sorted((atom_index, other_index)))
    raise RuntimeError("Could not locate an O-O bond in the ChemicalSystem")


def bond_distance_angstrom(system: ChemicalSystem, i: int, j: int) -> float:
    return float(system.get_distance(i, j))


def build_settings(
    atom_i: int,
    atom_j: int,
    start_distance: float,
    spinpolarization: int,
) -> Settings:
    settings = Settings()
    settings.runscript.nproc = 1
    settings.runscript.preamble_lines = ["export OMP_NUM_THREADS=1"]
    settings.input.ams.task = "PESScan"
    settings.input.ams.pesscan.scancoordinate = [Settings()]
    settings.input.ams.pesscan.scancoordinate[0].distance = [
        f"{atom_i + 1} {atom_j + 1} {start_distance:.6f} {TARGET_DISTANCE_ANGSTROM:.6f}"
    ]
    npoints = int(round((TARGET_DISTANCE_ANGSTROM - start_distance) / STEP_ANGSTROM)) + 1
    settings.input.ams.pesscan.scancoordinate[0].npoints = max(npoints, 2)
    settings.input.ams.pesscan.calcpropertiesatpespoints = "Yes"
    settings.input.mlpotential.model = MODEL
    settings.input.mlpotential.unrestricted = "Yes"
    settings.input.mlpotential.unpairedelectrons = spinpolarization
    return settings


def validate_settings(settings: Settings) -> None:
    AMS.from_settings(settings)


def run_scan(
    base_system: ChemicalSystem,
    atom_i: int,
    atom_j: int,
    spinpolarization: int,
) -> AMSJob:
    system = base_system.copy()
    start_distance = bond_distance_angstrom(system, atom_i, atom_j)
    settings = build_settings(atom_i, atom_j, start_distance, spinpolarization)
    validate_settings(settings)
    job = AMSJob(
        molecule=system,
        settings=settings,
        name=f"spinpol_{spinpolarization}",
    )
    result = job.run()
    if not result.ok():
        raise RuntimeError(f"Job {job.name} failed")
    return job


def main() -> None:
    init(folder=WORKDIR)
    try:
        system = build_system()
        atom_i, atom_j = find_oo_bond(system)
        start_distance = bond_distance_angstrom(system, atom_i, atom_j)
        print(f"O-O bond found between atoms {atom_i} and {atom_j} (0-based indexing)")
        print(f"O-O bond found between atoms {atom_i + 1} and {atom_j + 1} (AMS 1-based indexing)")
        print(f"Starting O-O distance: {start_distance:.3f} angstrom")

        for spinpolarization in (0, 1, 2):
            job = run_scan(system, atom_i, atom_j, spinpolarization)
            print(f"Finished {job.name}: {Path(job.path) / 'ams.rkf'}")
    finally:
        finish()


if __name__ == "__main__":
    main()