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

import math
from pathlib import Path

import numpy as np
from scm.base import ChemicalSystem, InputParser
from scm.plams import AMSJob, Settings, finish, init
from scm.utils.conversions import plams_molecule_to_chemsys


SMILES = "C=CC(CCCC)C1C=CC(C)=CC1=O"
SCAN_END_ANGSTROM = 1.5
SCAN_STEP_ANGSTROM = 0.2


def bonded_neighbors(system: ChemicalSystem, atom_index: int) -> list[tuple[int, float]]:
    """Return zero-based neighboring atom indices and bond orders."""
    neighbors: list[tuple[int, float]] = []
    for i, j, bond in system.bonds.get_bonds_for_atom(atom_index):
        other = j if i == atom_index else i
        neighbors.append((other, float(bond.order)))
    return neighbors


def find_reaction_atoms(system: ChemicalSystem) -> tuple[int, int]:
    """Find carbonyl O and terminal vinyl C from the ChemicalSystem bond graph."""
    carbonyl_oxygens: list[int] = []
    terminal_vinyl_carbons: list[int] = []
    for i, atom in enumerate(system):
        neighbors = bonded_neighbors(system, i)
        if atom.symbol == "O" and len(neighbors) == 1:
            j, order = neighbors[0]
            if system.atoms[j].symbol == "C" and order > 1.5:
                carbonyl_oxygens.append(i)
        if atom.symbol == "C":
            heavy_neighbors = [(j, order) for j, order in neighbors if system.atoms[j].symbol != "H"]
            carbon_double_bonds = [j for j, order in heavy_neighbors if system.atoms[j].symbol == "C" and order > 1.5]
            if len(heavy_neighbors) == 1 and len(carbon_double_bonds) == 1:
                terminal_vinyl_carbons.append(i)
    if len(carbonyl_oxygens) != 1 or len(terminal_vinyl_carbons) != 1:
        raise ValueError(
            "Expected one carbonyl oxygen and one terminal vinyl carbon; "
            f"found O={carbonyl_oxygens}, C={terminal_vinyl_carbons}"
        )
    return carbonyl_oxygens[0], terminal_vinyl_carbons[0]


def engine_settings() -> Settings:
    settings = Settings()
    settings.input.dftb.Model = "GFN1-xTB"
    return settings


def minimum_settings() -> Settings:
    settings = engine_settings()
    settings.input.ams.Task = "GeometryOptimization"
    settings.input.ams.Properties.NormalModes = "Yes"
    return settings


def scan_settings(oxygen: int, vinyl_carbon: int, start: float, npoints: int) -> Settings:
    settings = engine_settings()
    settings.input.ams.Task = "PESScan"
    settings.input.ams.PESScan.ScanCoordinate.nPoints = npoints
    settings.input.ams.PESScan.ScanCoordinate.Distance = (
        f"{oxygen + 1} {vinyl_carbon + 1} {start:.8f} {SCAN_END_ANGSTROM:.8f}"
    )
    return settings


def ts_settings(oxygen: int, vinyl_carbon: int) -> Settings:
    settings = engine_settings()
    settings.input.ams.Task = "TransitionStateSearch"
    settings.input.ams.Properties.NormalModes = "Yes"
    settings.input.ams.GeometryOptimization.InitialHessian.Type = "Calculate"
    settings.input.ams.TransitionStateSearch.ReactionCoordinate.Distance = (
        f"{oxygen + 1} {vinyl_carbon + 1} 1.0"
    )
    return settings


def validate(settings: Settings) -> None:
    input_text = AMSJob(settings=settings).get_input()
    InputParser().to_dict("ams", input_text)


def assert_success(job: AMSJob) -> None:
    if not job.ok():
        raise RuntimeError(f"AMS job failed: {job.path}")


def main() -> None:
    init(folder="01-run_workdir")
    initial = ChemicalSystem.from_smiles(SMILES, optimize_with_uff=True, num_trial_conformers=30)
    oxygen, vinyl_carbon = find_reaction_atoms(initial)
    print(f"Resolved carbonyl oxygen: atom {oxygen + 1}")
    print(f"Resolved terminal vinyl carbon: atom {vinyl_carbon + 1}")

    reactant_settings = minimum_settings()
    validate(reactant_settings)
    reactant_job = AMSJob(
        molecule=initial,
        settings=reactant_settings,
        name="reactant_opt_freq",
    )
    reactant_job.run()
    assert_success(reactant_job)
    reactant = reactant_job.results.get_main_system()

    optimized_distance = float(reactant.get_distance(oxygen, vinyl_carbon))
    n_intervals = int(math.ceil((optimized_distance - SCAN_END_ANGSTROM) / SCAN_STEP_ANGSTROM))
    scan_start = SCAN_END_ANGSTROM + n_intervals * SCAN_STEP_ANGSTROM
    npoints = n_intervals + 1
    print(f"Optimized reactant O...C distance: {optimized_distance:.8f} angstrom")
    print(f"Scan: {scan_start:.8f} to {SCAN_END_ANGSTROM:.8f} angstrom, {npoints} points")

    pes_settings = scan_settings(oxygen, vinyl_carbon, scan_start, npoints)
    validate(pes_settings)
    scan_job = AMSJob(molecule=reactant, settings=pes_settings, name="oc_distance_scan")
    scan_job.run()
    assert_success(scan_job)
    scan_results = scan_job.results.get_pesscan_results()
    energies = np.asarray(scan_results["PES"], dtype=float)
    converged = np.asarray(scan_results["Converged"], dtype=bool)
    if not np.all(converged):
        failed = (np.flatnonzero(~converged) + 1).tolist()
        raise RuntimeError(f"Unconverged PES scan points: {failed}")

    highest_index = int(np.argmax(energies))
    highest_system = plams_molecule_to_chemsys(scan_results["Molecules"][highest_index])
    print(f"Highest-energy scan point: {highest_index + 1} of {len(energies)}")

    transition_settings = ts_settings(oxygen, vinyl_carbon)
    validate(transition_settings)
    ts_job = AMSJob(molecule=highest_system, settings=transition_settings, name="ts_opt_freq")
    ts_job.run()
    assert_success(ts_job)

    product_start = plams_molecule_to_chemsys(scan_results["Molecules"][-1])
    product_settings = minimum_settings()
    validate(product_settings)
    product_job = AMSJob(
        molecule=product_start,
        settings=product_settings,
        name="product_opt_freq",
    )
    product_job.run()
    assert_success(product_job)

    ts_frequencies = np.asarray(ts_job.results.get_frequencies(unit="cm^-1"), dtype=float)
    negative = ts_frequencies[ts_frequencies < 0.0]
    character = ts_job.results.readrkf("AMSResults", "PESPointCharacter", file="engine")
    print(f"TS PES point character: {character}")
    print(f"TS negative frequencies (cm^-1): {negative.tolist()}")
    if len(negative) != 1:
        raise RuntimeError(f"TS verification failed: expected exactly one negative frequency, found {len(negative)}")

    Path("run_complete.txt").write_text(
        f"oxygen_atom_1based={oxygen + 1}\n"
        f"terminal_vinyl_carbon_1based={vinyl_carbon + 1}\n"
        f"optimized_reactant_distance_angstrom={optimized_distance:.8f}\n"
        f"scan_start_angstrom={scan_start:.8f}\n"
        f"scan_end_angstrom={SCAN_END_ANGSTROM:.8f}\n"
        f"scan_step_angstrom={SCAN_STEP_ANGSTROM:.8f}\n"
        f"scan_points={npoints}\n"
        f"highest_scan_point_1based={highest_index + 1}\n"
        f"ts_negative_frequencies_cm-1={negative.tolist()}\n"
        f"ts_pes_point_character={character}\n"
    )
    finish()


if __name__ == "__main__":
    main()