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

import os
import time
from pathlib import Path
from typing import Any

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

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 adf_settings(task: str, normal_modes: bool = False, hessian_region: str | None = None, vib: bool = False) -> Settings:
    s = Settings()
    s.input.ams.Task = task
    if normal_modes:
        s.input.ams.Properties.NormalModes = "Yes"
        s.input.ams.NormalModes.ReScanModes = "No"
    if hessian_region is not None:
        s.input.ams.Properties.SelectedRegionForHessian = hessian_region
    if vib:
        s.input.ams.Task = "VibrationalAnalysis"
    s.input.adf.XC.GGA = "PBE"
    s.input.adf.Basis.Type = "DZ"
    s.input.adf.Basis.Core = "Large"
    # Keep defaults otherwise, per requested method.
    AMS.from_settings(s)
    return s


def find_nonring_carbonyl(system: ChemicalSystem) -> tuple[int, int]:
    candidates: list[tuple[int, int, float]] = []
    for i, atom in enumerate(system.atoms):
        if atom.symbol != "C":
            continue
        for _frm, j, bond in system.bonds.get_bonds_for_atom(i):
            if j <= i:
                continue
            other = system.atoms[j]
            if other.symbol != "O":
                continue
            order = getattr(bond, "order", None)
            if order is None:
                order = getattr(bond, "bond_order", 0.0)
            if float(order) > 1.5 and (not system.atom_is_in_ring(i)):
                candidates.append((i, j, float(order)))
    if len(candidates) != 1:
        raise RuntimeError(f"Expected one non-ring C=O; found {candidates}")
    return candidates[0][0], candidates[0][1]


def mode_purity(job: AMSJob, cidx: int, oidx: int) -> tuple[int, float, float, np.ndarray]:
    freqs = np.array(job.results.get_frequencies())
    modes = np.array(job.results.get_normal_modes())
    n_atoms = len(job.results.get_main_system())
    # modes shape is often (nModes, nAtoms, 3); handle transposed variants.
    if modes.shape[1] != n_atoms:
        if modes.shape[-1] == n_atoms:
            modes = np.moveaxis(modes, -1, 1)
    scores = []
    for m in modes:
        atom_sq = np.sum(m * m, axis=1)
        co = float(atom_sq[cidx] + atom_sq[oidx])
        total = float(np.sum(atom_sq))
        scores.append(co / total if total else 0.0)
    best0 = int(np.argmax(scores))
    return best0 + 1, float(freqs[best0]), float(scores[best0]), modes[best0]


def main() -> None:
    init(folder="01-run_workdir")
    system = ChemicalSystem.from_smiles(SMILES)
    system.guess_bonds()
    cidx, oidx = find_nonring_carbonyl(system)
    system.add_atom_to_region(cidx, "CO")
    system.add_atom_to_region(oidx, "CO")
    print(f"Selected non-ring carbonyl (1-based): C={cidx+1}, O={oidx+1}")

    jobs: list[AMSJob] = []
    timings: dict[str, float] = {}

    # 1 Geometry optimization
    t0 = time.perf_counter()
    opt = AMSJob(settings=adf_settings("GeometryOptimization"), molecule=system, name="01_opt")
    opt.run()
    timings["01_opt"] = time.perf_counter() - t0
    if not opt.ok():
        raise RuntimeError("Geometry optimization failed")
    jobs.append(opt)

    opt_system = opt.results.get_main_system()
    # Preserve/restore the CO region on optimized system, based on original indices.
    opt_system.add_atom_to_region(cidx, "CO")
    opt_system.add_atom_to_region(oidx, "CO")

    # 2 partial Hessian / normal modes on CO region
    t0 = time.perf_counter()
    ph = AMSJob(settings=adf_settings("SinglePoint", normal_modes=True, hessian_region="CO"), molecule=opt_system, name="02_partial_hessian")
    ph.run()
    timings["02_partial_hessian"] = time.perf_counter() - t0
    if not ph.ok():
        raise RuntimeError("Partial Hessian failed")
    jobs.append(ph)
    mode_no, f, purity, _mode = mode_purity(ph, cidx, oidx)
    print(f"Selected partial-Hessian mode {mode_no}: {f:.2f} cm^-1, CO displacement fraction {purity:.3f}")

    # 3 mode tracking from selected partial-Hessian mode
    mt_s = adf_settings("VibrationalAnalysis")
    mt_s.input.ams.VibrationalAnalysis.Type = "ModeTracking"
    mt_s.input.ams.VibrationalAnalysis.NormalModes.ModeInputFormat = "File"
    mt_s.input.ams.VibrationalAnalysis.NormalModes.ModeFile = str(Path(ph.path) / "adf.rkf")
    mt_s.input.ams.VibrationalAnalysis.NormalModes.ModeSelect.ModeNumber = [mode_no]
    mt_s.input.ams.VibrationalAnalysis.ModeTracking.HessianGuess = "File"
    mt_s.input.ams.VibrationalAnalysis.ModeTracking.HessianPath = str(Path(ph.path))
    mt_s.input.ams.VibrationalAnalysis.ModeTracking.TrackingMethod = "OverlapPrevious"
    AMS.from_settings(mt_s)
    t0 = time.perf_counter()
    mt = AMSJob(settings=mt_s, molecule=opt_system, name="03_mode_tracking")
    mt.run()
    timings["03_mode_tracking"] = time.perf_counter() - t0
    if not mt.ok():
        raise RuntimeError("Mode tracking failed")
    jobs.append(mt)

    # 4 full normal modes verification
    t0 = time.perf_counter()
    full = AMSJob(settings=adf_settings("SinglePoint", normal_modes=True), molecule=opt_system, name="04_full_normal_modes")
    full.run()
    timings["04_full_normal_modes"] = time.perf_counter() - t0
    if not full.ok():
        raise RuntimeError("Full normal modes failed")
    jobs.append(full)

    with open("timings.txt", "w") as fh:
        fh.write(f"carbonyl_C_1based {cidx+1}\ncarbonyl_O_1based {oidx+1}\npartial_mode {mode_no}\n")
        for k, v in timings.items():
            fh.write(f"{k} {v:.3f}\n")
    finish()

if __name__ == "__main__":
    main()