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

import time
from pathlib import Path

from scm.input_classes import AMS
from scm.plams import AMSJob, init, finish

from importlib.machinery import SourceFileLoader
mod = SourceFileLoader("runmod", "01-run.py").load_module()


def main() -> None:
    init(folder="02-continue_workdir")
    cidx, oidx = 1, 2  # zero-based, found from ChemicalSystem bonds in 01-run.py: C=2, O=3 (1-based)
    opt = AMSJob.load_external("01-run_workdir/01_opt")
    opt_system = opt.results.get_main_system()
    opt_system.add_atom_to_region(cidx, "CO")
    opt_system.add_atom_to_region(oidx, "CO")
    timings: dict[str, float] = {}

    t0 = time.perf_counter()
    ph = AMSJob(settings=mod.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")
    mode_no, f, purity, _mode = mod.mode_purity(ph, cidx, oidx)
    print(f"Selected partial-Hessian mode {mode_no}: {f:.2f} cm^-1, CO displacement fraction {purity:.3f}")

    mt_s = mod.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")

    t0 = time.perf_counter()
    full = AMSJob(settings=mod.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")

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

if __name__ == "__main__":
    main()