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

from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scm.plams import AMSJob, plot_image_grid, view


ROOT = Path(__file__).resolve().parent
JOB_NAMES = [
    "01_dftb_preopt",
    "02_dft_opt",
    "03_dft_partial_hessian",
    "04_dftb_hessian_guess",
    "05_dft_mode_tracking",
    "06_dft_full_modes",
]
JUSTIFICATIONS = {
    "01_dftb_preopt": "GFN1-xTB preoptimization and initial Hessian for the DFT optimization.",
    "02_dft_opt": "Reference PBE/DZ geometry and ADF engine restart data.",
    "03_dft_partial_hessian": "Two-atom CO partial Hessian used to construct the tracked-mode guess.",
    "04_dftb_hessian_guess": "GFN1-xTB full Hessian at the DFT geometry, used as the mode-tracking preconditioner.",
    "05_dft_mode_tracking": "PBE/DZ mode-tracking result for the selected CO-stretch guess.",
    "06_dft_full_modes": "PBE/DZ analytical full-Hessian reference calculation.",
}


def latest_complete_workdir() -> Path:
    candidates = sorted(
        ROOT.glob("01-run_workdir*"), key=lambda path: path.stat().st_mtime, reverse=True
    )
    for candidate in candidates:
        if all((candidate / name).is_dir() for name in JOB_NAMES):
            return candidate
    raise FileNotFoundError("No complete 01-run_workdir job set found")


def load_jobs(workdir: Path) -> dict[str, AMSJob]:
    jobs: dict[str, AMSJob] = {}
    for name in JOB_NAMES:
        # Each load is paired with JUSTIFICATIONS[name] in the rendered provenance section.
        jobs[name] = AMSJob.load_external(str(workdir / name))
    return jobs


def selected_partial_mode(partial: AMSJob, carbon_index: int, oxygen_index: int) -> int:
    modes = np.asarray(partial.results.get_normal_modes())
    coords = np.asarray(partial.results.get_main_system().coords)
    bond = coords[oxygen_index] - coords[carbon_index]
    bond /= np.linalg.norm(bond)
    scores: list[float] = []
    for mode in modes:
        denominator = np.linalg.norm(mode[[carbon_index, oxygen_index]])
        relative = mode[oxygen_index] - mode[carbon_index]
        score = 0.0 if denominator < 1.0e-14 else abs(float(relative @ bond)) / denominator
        scores.append(score)
    return int(np.argmax(scores))


def normalized_overlap(left: np.ndarray, right: np.ndarray) -> float:
    left_flat = left.reshape(-1)
    right_flat = right.reshape(-1)
    return abs(float(left_flat @ right_flat)) / (
        float(np.linalg.norm(left_flat)) * float(np.linalg.norm(right_flat))
    )


def spectrum_figure(job: AMSJob, title: str, path: Path, engine: str | None = None) -> None:
    x_values, y_values = job.results.get_ir_spectrum(
        engine=engine,
        broadening_type="gaussian",
        broadening_width=20,
        min_x=0,
        max_x=4000,
        x_spacing=1.0,
    )
    fig, axis = plt.subplots(figsize=(9, 4.2))
    axis.plot(x_values, y_values, color="#204a87", linewidth=1.2)
    axis.set(xlabel="Wavenumber (cm$^{-1}$)", ylabel="IR intensity (km mol$^{-1}$)", title=title)
    axis.set_xlim(4000, 0)
    axis.grid(alpha=0.2)
    fig.tight_layout()
    fig.savefig(path, dpi=180)
    plt.close(fig)


def elapsed_seconds(job: AMSJob) -> float:
    timings = job.results.get_timings()
    lowered = {str(key).lower(): float(value) for key, value in timings.items()}
    for key in ("elapsed", "wall", "total"):
        if key in lowered:
            return lowered[key]
    return max(lowered.values())


def main() -> None:
    workdir = latest_complete_workdir()
    jobs = load_jobs(workdir)
    figures = ROOT / "figures"
    tables = ROOT / "tables"
    figures.mkdir(exist_ok=True)
    tables.mkdir(exist_ok=True)

    optimized = jobs["02_dft_opt"].results.get_main_system()
    carbon_index, oxygen_index = tuple(
        int(index) for index in optimized.get_atoms_in_region("CO")
    )
    if optimized.atoms[carbon_index].symbol == "O":
        carbon_index, oxygen_index = oxygen_index, carbon_index

    partial = jobs["03_dft_partial_hessian"]
    tracking = jobs["05_dft_mode_tracking"]
    full = jobs["06_dft_full_modes"]
    partial_index = selected_partial_mode(partial, carbon_index, oxygen_index)

    partial_frequencies = np.asarray(partial.results.get_frequencies(unit="cm^-1"))
    partial_intensities = np.asarray(partial.results.get_ir_intensities())
    tracking_frequencies = np.asarray(
        tracking.results.get_frequencies(unit="cm^-1", engine="ams")
    )
    tracking_intensities = np.asarray(tracking.results.get_ir_intensities(engine="ams"))
    tracking_modes = np.asarray(
        tracking.results.readrkf("Vibrations", "NoWeightNormalMode(1)", file="ams")
    ).reshape(1, -1, 3)
    full_frequencies = np.asarray(full.results.get_frequencies(unit="cm^-1"))
    full_intensities = np.asarray(full.results.get_ir_intensities())
    full_modes = np.asarray(full.results.get_normal_modes())

    tracked_mode = tracking_modes[0]
    overlaps = np.array([normalized_overlap(tracked_mode, mode) for mode in full_modes])
    full_index = int(np.argmax(overlaps))

    mode_table = pd.DataFrame(
        [
            {
                "Calculation": "Partial Hessian guess",
                "Mode number": partial_index + 1,
                "Frequency (cm^-1)": partial_frequencies[partial_index],
                "IR intensity (km/mol)": partial_intensities[partial_index],
                "Overlap with tracked mode": np.nan,
            },
            {
                "Calculation": "Mode tracking",
                "Mode number": 1,
                "Frequency (cm^-1)": tracking_frequencies[0],
                "IR intensity (km/mol)": tracking_intensities[0],
                "Overlap with tracked mode": 1.0,
            },
            {
                "Calculation": "Full normal modes",
                "Mode number": full_index + 1,
                "Frequency (cm^-1)": full_frequencies[full_index],
                "IR intensity (km/mol)": full_intensities[full_index],
                "Overlap with tracked mode": overlaps[full_index],
            },
        ]
    )
    mode_table.to_csv(tables / "selected_mode.csv", index=False)

    timing_rows = []
    labels = {
        "01_dftb_preopt": "DFTB preoptimization + modes",
        "02_dft_opt": "DFT geometry optimization",
        "03_dft_partial_hessian": "DFT partial Hessian",
        "04_dftb_hessian_guess": "DFTB Hessian guess",
        "05_dft_mode_tracking": "DFT mode tracking",
        "06_dft_full_modes": "DFT full normal modes",
    }
    for name in JOB_NAMES:
        timing_rows.append({"Job": labels[name], "Elapsed time (s)": elapsed_seconds(jobs[name])})
    partial_tracking_time = elapsed_seconds(partial) + elapsed_seconds(tracking)
    full_time = elapsed_seconds(full)
    timing_rows.extend(
        [
            {"Job": "DFT partial Hessian + DFT mode tracking", "Elapsed time (s)": partial_tracking_time},
            {"Job": "Full DFT normal modes", "Elapsed time (s)": full_time},
        ]
    )
    timing_table = pd.DataFrame(timing_rows)
    timing_table["Elapsed time (min)"] = timing_table["Elapsed time (s)"] / 60.0
    timing_table.to_csv(tables / "timings.csv", index=False)

    spectrum_figure(partial, "Partial-Hessian IR spectrum", figures / "partial_hessian_spectrum.png")
    spectrum_figure(
        tracking,
        "Mode-tracking IR spectrum",
        figures / "mode_tracking_spectrum.png",
        engine="ams",
    )
    spectrum_figure(full, "Full analytical-Hessian IR spectrum", figures / "full_spectrum.png")

    scale = 0.55 / float(np.max(np.linalg.norm(tracked_mode, axis=1)))
    negative = optimized.copy()
    positive = optimized.copy()
    negative.coords = np.asarray(optimized.coords) - scale * tracked_mode
    positive.coords = np.asarray(optimized.coords) + scale * tracked_mode
    images = {
        "Mode displaced, negative": view(
            negative,
            guess_bonds=len(negative.bonds) == 0,
            direction="along_pca3",
            show_regions=True,
            width=500,
            height=420,
            picture_path=str(figures / "mode_negative.png"),
        ),
        "Optimized geometry": view(
            optimized,
            guess_bonds=len(optimized.bonds) == 0,
            direction="along_pca3",
            show_regions=True,
            width=500,
            height=420,
            picture_path=str(figures / "optimized_geometry.png"),
        ),
        "Mode displaced, positive": view(
            positive,
            guess_bonds=len(positive.bonds) == 0,
            direction="along_pca3",
            show_regions=True,
            width=500,
            height=420,
            picture_path=str(figures / "mode_positive.png"),
        ),
    }
    plot_image_grid(images, rows=1, save_path=str(figures / "mode_displacement_grid.png"))

    frequency_error = abs(float(tracking_frequencies[0] - full_frequencies[full_index]))
    frequency_relative_error = 100.0 * frequency_error / abs(float(full_frequencies[full_index]))
    intensity_error = abs(float(tracking_intensities[0] - full_intensities[full_index]))
    intensity_relative_error = (
        100.0 * intensity_error / abs(float(full_intensities[full_index]))
        if abs(float(full_intensities[full_index])) > 1.0e-14
        else float("nan")
    )
    speedup = full_time / partial_tracking_time
    focused_with_guess_time = partial_tracking_time + elapsed_seconds(jobs["04_dftb_hessian_guess"])
    inclusive_speedup = full_time / focused_with_guess_time

    lines = [
        "# Mode tracking versus a full normal-mode calculation",
        "",
        "## Introduction",
        "",
        "Mode tracking calculates one chosen normal mode without building the complete molecular Hessian. It starts from a physically motivated displacement and iteratively improves that vector until it is an eigenvector of the target-method Hessian. This report tests whether that focused calculation reproduces the PBE/DZ analytical normal mode for the non-ring C=O stretch of dydrogesterone, and measures the time saved against a full normal-mode calculation.",
        "",
        "Dydrogesterone was built from its stereochemical SMILES. Bond data identified atom " + str(oxygen_index + 1) + " as the O double-bonded to non-ring atom " + str(carbon_index + 1) + " C. These two atoms form region `CO`. The geometry was preoptimized with GFN1-xTB and normal modes, then optimized with ADF/PBE/DZ, large frozen core, and NumericalQuality Normal. The DFT optimization used the GFN1-xTB Hessian as its initial Hessian. Every subsequent calculation used the final DFT geometry without further optimization.",
        "",
        "## Tracked structure and displacement",
        "",
        "The translucent region marks the selected carbonyl atoms. The outer structures show equal negative and positive displacements along the converged tracked mode.",
        "",
        "![Negative displacement, optimized geometry, and positive displacement](figures/mode_displacement_grid.png)",
        "",
        "## Spectra",
        "",
        "### Partial Hessian",
        "",
        "![Partial-Hessian spectrum](figures/partial_hessian_spectrum.png)",
        "",
        "### Mode tracking",
        "",
        "![Mode-tracking spectrum](figures/mode_tracking_spectrum.png)",
        "",
        "### Full normal modes",
        "",
        "![Full normal-mode spectrum](figures/full_spectrum.png)",
        "",
        "## Selected-mode accuracy",
        "",
        mode_table.to_markdown(index=False, floatfmt=".4f"),
        "",
        f"The full-spectrum mode was matched by normal-mode overlap, not frequency proximity. Its absolute overlap with the tracked vector is {overlaps[full_index]:.6f}. Mode tracking differs from the full calculation by {frequency_error:.4f} cm^-1 in frequency, or {frequency_relative_error:.3f}%, and {intensity_error:.4f} km/mol in IR intensity, or {intensity_relative_error:.3f}%.",
        "",
        "## Timings",
        "",
        timing_table.to_markdown(index=False, floatfmt=".3f"),
        "",
        f"The two DFT steps specific to the focused route, partial Hessian plus mode tracking, took {partial_tracking_time:.3f} s. The full analytical normal-mode job took {full_time:.3f} s, so the DFT-only focused route was {speedup:.2f}x faster. Including the {elapsed_seconds(jobs['04_dftb_hessian_guess']):.3f} s GFN1-xTB Hessian-guess job gives a focused-route time of {focused_with_guess_time:.3f} s and a {inclusive_speedup:.2f}x speedup. Geometry optimization timings are excluded from both comparisons because both vibrational routes use the same optimized geometry.",
        "",
        "## Conclusion",
        "",
        f"For this carbonyl mode, mode tracking reproduced the full-Hessian reference with a frequency error of {frequency_error:.4f} cm^-1, an intensity error of {intensity_error:.4f} km/mol, and a mode-vector overlap of {overlaps[full_index]:.6f}. The full calculation had no imaginary frequencies. The DFT partial-Hessian plus mode-tracking steps were {speedup:.2f}x faster than the full analytical Hessian, or {inclusive_speedup:.2f}x faster after including the DFTB Hessian guess.",
        "",
        "## Calculation provenance and inputs",
        "",
        f"Jobs were loaded directly from `{workdir.name}`.",
        "",
    ]
    for name in JOB_NAMES:
        lines.extend(
            [
                f"### {labels[name]}",
                "",
                JUSTIFICATIONS[name],
                "",
                "```ams",
                jobs[name].get_input().rstrip(),
                "```",
                "",
            ]
        )
    lines.extend(
        [
            "## References",
            "",
            "- Amsterdam Modeling Suite 2026 documentation, `Doc/text/Tutorials/VibrationalSpectroscopy/ModeTracking.txt`, and the AMS Vibrational Spectroscopy manual.",
            "",
            "- S. Luber, J. Neugebauer, and M. Reiher, \"Intensity tracking for theoretical infrared spectroscopy of large molecules,\" *Journal of Chemical Physics* **130**, 064105 (2009).",
            "",
            "- G. L. G. Sleijpen and H. A. van der Vorst, \"A Jacobi-Davidson iteration method for linear eigenvalue problems,\" *SIAM Journal on Matrix Analysis and Applications* **17**, 401 (1996).",
            "",
        ]
    )
    (ROOT / "report.md").write_text("\n".join(lines), encoding="utf-8")
    print(f"Wrote {ROOT / 'report.md'}")


if __name__ == "__main__":
    main()