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

import time
from pathlib import Path
import numpy as np
from importlib.machinery import SourceFileLoader
from scm.input_classes import AMS
from scm.plams import AMSJob, init, finish

mod = SourceFileLoader("runmod", "01-run.py").load_module()


def main() -> None:
    init(folder="03-rest_workdir")
    cidx, oidx = 1, 2
    opt = AMSJob.load_external("01-run_workdir/01_opt")
    ph = AMSJob.load_external("02-continue_workdir/02_partial_hessian")
    opt_system = opt.results.get_main_system()
    opt_system.add_atom_to_region(cidx, "CO")
    opt_system.add_atom_to_region(oidx, "CO")
    freqs = np.array(ph.results.get_frequencies())
    nonzero = np.where(np.abs(freqs) > 1.0)[0]
    mode_no = int(nonzero[np.argmax(freqs[nonzero])] + 1)
    print(f"Selected highest non-zero partial-Hessian C=O stretch mode {mode_no}: {freqs[mode_no-1]:.2f} cm^-1")
    timings: dict[str, float] = {}

    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) / "adf.rkf")
    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()