#!/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, init


SMILES = "c1ccccc1-c2ccccc2"
JOB_NAME = "biphenyl_dftb_torsion"


@dataclass(frozen=True)
class TorsionDefinition:
    central_bond: tuple[int, int]
    dihedral: tuple[int, int, int, int]

    @property
    def ams_dihedral_line(self) -> str:
        a, b, c, d = (idx + 1 for idx in self.dihedral)
        return f"{a} {b} {c} {d} 0.0 180.0"


def carbon_neighbors(system: ChemicalSystem, atom_index: int, exclude: int) -> list[int]:
    return [
        idx
        for idx in system.bonds.get_bonded_atoms(atom_index)
        if idx != exclude and system.atoms[idx].symbol == "C"
    ]


def identify_inter_ring_torsion(system: ChemicalSystem) -> TorsionDefinition:
    candidates: list[tuple[int, int]] = []
    for i, j, _bond in system.bonds:
        if system.atoms[i].symbol != "C" or system.atoms[j].symbol != "C":
            continue
        if system.bond_cuts_molecule(i, j):
            candidates.append((i, j))

    if len(candidates) != 1:
        raise RuntimeError(f"Expected one central inter-ring C-C bond, found {candidates}")

    left, right = candidates[0]
    left_neighbors = carbon_neighbors(system, left, exclude=right)
    right_neighbors = carbon_neighbors(system, right, exclude=left)
    if not left_neighbors or not right_neighbors:
        raise RuntimeError("Could not find ring-neighbor carbon atoms for the inter-ring dihedral")

    return TorsionDefinition(
        central_bond=(left, right),
        dihedral=(left_neighbors[0], left, right, right_neighbors[0]),
    )


def make_settings(torsion: TorsionDefinition) -> Settings:
    settings = Settings()
    settings.input.ams.task = "PESScan"
    settings.input.ams.pesscan.optimize = True
    settings.input.ams.pesscan.scancoordinate.npoints = 19
    settings.input.ams.pesscan.scancoordinate.dihedral = torsion.ams_dihedral_line
    settings.input.dftb.model = "GFN1-xTB"

    AMS.from_settings(settings)
    return settings


def main() -> None:
    init(folder="01-run-biphenyl-dftb-pesscan_workdir")

    system = ChemicalSystem.from_smiles(SMILES)
    torsion = identify_inter_ring_torsion(system)
    settings = make_settings(torsion)

    print(f"Built {system.formula()} from SMILES: {SMILES}")
    print(f"Central inter-ring C-C bond, zero-based: {torsion.central_bond}")
    print(f"Inter-ring dihedral, zero-based: {torsion.dihedral}")
    print(f"AMS dihedral scan line: {torsion.ams_dihedral_line}")

    job = AMSJob(molecule=system, settings=settings, name=JOB_NAME)
    result = job.run()
    if not result.ok():
        raise RuntimeError(f"AMS job failed; inspect {job.path}")


if __name__ == "__main__":
    main()