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

import numpy as np
from scm.base import ChemicalSystem, InputParser
from scm.plams import AMSJob, Settings, init


SMILES = "CC(=O)[C@H]1CC[C@@H]2[C@@]1(CC[C@@H]3[C@H]2C=CC4=CC(=O)CC[C@@]34C)C"


def identify_nonring_carbonyl(system: ChemicalSystem) -> tuple[int, int]:
    """Return zero-based C and O indices for the non-ring carbonyl."""
    matches: list[tuple[int, int]] = []
    for oxygen_index, oxygen in enumerate(system.atoms):
        if oxygen.symbol != "O":
            continue
        for carbon_index in system.bonds.get_bonded_atoms(oxygen_index):
            if system.atoms[carbon_index].symbol != "C":
                continue
            bonds = list(system.bonds.get_bonds_between_atoms(oxygen_index, carbon_index))
            if any(abs(bond.order - 2.0) < 1.0e-8 for _, _, bond in bonds):
                if not system.atom_is_in_ring(carbon_index):
                    matches.append((carbon_index, oxygen_index))
    if len(matches) != 1:
        raise RuntimeError(f"Expected one non-ring carbonyl, found {matches}")
    return matches[0]


def dft_settings() -> Settings:
    settings = Settings()
    settings.input.adf.XC.GGA = "PBE"
    settings.input.adf.Basis.Type = "DZ"
    settings.input.adf.Basis.Core = "Large"
    settings.input.adf.NumericalQuality = "Normal"
    return settings


def validate(job: AMSJob) -> None:
    InputParser().to_dict("ams", job.get_input())


def run_checked(job: AMSJob) -> AMSJob:
    validate(job)
    print(f"\nValidated input for {job.name}\n", flush=True)
    result = job.run()
    if not result.ok():
        raise RuntimeError(f"AMS job failed: {job.name}")
    timings = job.results.get_timings()
    print(f"Completed {job.name}; timings={timings}", flush=True)
    return job


def restarted_adf_settings(engine_path: str) -> Settings:
    settings = Settings()
    settings.input.ams.LoadEngine = engine_path
    settings.input.ams.EngineRestart = engine_path
    return settings


def stretch_score(
    mode: np.ndarray, coords: np.ndarray, carbon_index: int, oxygen_index: int
) -> float:
    bond_vector = coords[oxygen_index] - coords[carbon_index]
    bond_unit = bond_vector / np.linalg.norm(bond_vector)
    relative_displacement = mode[oxygen_index] - mode[carbon_index]
    denominator = np.linalg.norm(mode[[carbon_index, oxygen_index]])
    if denominator < 1.0e-14:
        return 0.0
    return abs(float(np.dot(relative_displacement, bond_unit))) / denominator


def main() -> None:
    init(folder="01-run_workdir")

    system = ChemicalSystem.from_smiles(SMILES)
    carbon_index, oxygen_index = identify_nonring_carbonyl(system)
    system.set_atoms_in_region(np.array([carbon_index, oxygen_index]), "CO")
    print(
        f"Non-ring carbonyl: C atom {carbon_index + 1}, O atom {oxygen_index + 1}; "
        "both assigned to region CO",
        flush=True,
    )

    preopt_settings = Settings()
    preopt_settings.input.ams.Task = "GeometryOptimization"
    preopt_settings.input.ams.Properties.NormalModes = "Yes"
    preopt_settings.input.dftb.Model = "GFN1-xTB"
    preopt = run_checked(
        AMSJob(molecule=system, settings=preopt_settings, name="01_dftb_preopt")
    )

    preoptimized = preopt.results.get_main_system()
    dft_opt_settings = dft_settings()
    dft_opt_settings.input.ams.Task = "GeometryOptimization"
    dft_opt_settings.input.ams.GeometryOptimization.InitialHessian.Type = "FromFile"
    dft_opt_settings.input.ams.GeometryOptimization.InitialHessian.File = preopt.results.rkfpath(
        file="engine"
    )
    dft_opt = run_checked(
        AMSJob(molecule=preoptimized, settings=dft_opt_settings, name="02_dft_opt")
    )

    optimized = dft_opt.results.get_main_system()
    dft_engine_path = dft_opt.results.rkfpath(file="engine")

    partial_settings = restarted_adf_settings(dft_engine_path)
    partial_settings.input.ams.Task = "SinglePoint"
    partial_settings.input.ams.Properties.NormalModes = "Yes"
    partial_settings.input.ams.Properties.SelectedRegionForHessian = "CO"
    partial = run_checked(
        AMSJob(molecule=optimized, settings=partial_settings, name="03_dft_partial_hessian")
    )

    partial_modes = np.asarray(partial.results.get_normal_modes())
    partial_frequencies = np.asarray(partial.results.get_frequencies(unit="cm^-1"))
    coords = np.asarray(optimized.coords)
    scores = np.array(
        [
            stretch_score(mode, coords, carbon_index, oxygen_index)
            for mode in partial_modes
        ]
    )
    mode_index = int(np.argmax(scores))
    mode_number = mode_index + 1
    print(
        f"Selected partial-Hessian mode {mode_number}: "
        f"frequency={partial_frequencies[mode_index]:.6f} cm^-1, "
        f"normalized C=O stretch score={scores[mode_index]:.8f}",
        flush=True,
    )
    top = np.argsort(scores)[-10:][::-1]
    for idx in top:
        print(
            f"  candidate mode {idx + 1:3d}: {partial_frequencies[idx]:12.5f} cm^-1, "
            f"score={scores[idx]:.8f}",
            flush=True,
        )

    dftb_freq_settings = Settings()
    dftb_freq_settings.input.ams.Task = "SinglePoint"
    dftb_freq_settings.input.ams.Properties.NormalModes = "Yes"
    dftb_freq_settings.input.dftb.Model = "GFN1-xTB"
    dftb_freq = run_checked(
        AMSJob(molecule=optimized, settings=dftb_freq_settings, name="04_dftb_hessian_guess")
    )

    tracking_settings = restarted_adf_settings(dft_engine_path)
    tracking_settings.input.ams.Task = "VibrationalAnalysis"
    tracking_settings.input.ams.VibrationalAnalysis.Type = "ModeTracking"
    tracking_settings.input.ams.VibrationalAnalysis.NormalModes.ModeInputFormat = "File"
    tracking_settings.input.ams.VibrationalAnalysis.NormalModes.ModeFile = partial.results.rkfpath(
        file="engine"
    )
    tracking_settings.input.ams.VibrationalAnalysis.NormalModes.ModeSelect.ModeNumber = str(
        mode_number
    )
    tracking_settings.input.ams.VibrationalAnalysis.ModeTracking.HessianGuess = "File"
    tracking_settings.input.ams.VibrationalAnalysis.ModeTracking.HessianPath = (
        dftb_freq.results.rkfpath(file="engine")
    )
    tracking = run_checked(
        AMSJob(molecule=optimized, settings=tracking_settings, name="05_dft_mode_tracking")
    )

    full_settings = restarted_adf_settings(dft_engine_path)
    full_settings.input.ams.Task = "SinglePoint"
    full_settings.input.ams.Properties.NormalModes = "Yes"
    run_checked(AMSJob(molecule=optimized, settings=full_settings, name="06_dft_full_modes"))

    print("All calculations completed successfully.", flush=True)


if __name__ == "__main__":
    main()