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

import re
from datetime import datetime
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scm.base import ChemicalSystem
from scm.plams import AMSJob, init, finish, view, plot_image_grid

CIDX = 1  # zero-based (1-based: 2)
OIDX = 2  # zero-based (1-based: 3)
PARTIAL_MODE = 153

JOBS = {
    "opt": ("01-run_workdir/01_opt", "Geometry optimization at ADF/PBE/DZ large core."),
    "partial": ("02-continue_workdir/02_partial_hessian", "Partial Hessian / normal-modes calculation restricted to CO region."),
    "mt": ("03-rest_workdir/03_mode_tracking", "Mode tracking started from partial-Hessian C=O stretch mode 153."),
    "full": ("03-rest_workdir/04_full_normal_modes", "Full normal-modes verification calculation."),
}


def load_jobs() -> dict[str, AMSJob]:
    return {k: AMSJob.load_external(path) for k, (path, _why) in JOBS.items()}


def vib_arrays(job: AMSJob, file: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    f = np.atleast_1d(job.results.readrkf("Vibrations", "Frequencies[cm-1]", file=file)).astype(float)
    inten = np.atleast_1d(job.results.readrkf("Vibrations", "Intensities[km/mol]", file=file)).astype(float)
    modes = []
    for i in range(1, len(f) + 1):
        arr = np.array(job.results.readrkf("Vibrations", f"NoWeightNormalMode({i})", file=file), dtype=float)
        arr = arr.reshape(-1, 3)
        modes.append(arr)
    return f, inten, np.array(modes)


def co_fraction(mode: np.ndarray) -> float:
    atom_sq = np.sum(mode * mode, axis=1)
    return float((atom_sq[CIDX] + atom_sq[OIDX]) / np.sum(atom_sq))


def normalize(mode: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(mode, axis=1)
    m = np.max(norms)
    return mode / m if m else mode


def parse_duration(logfile: str, jobname: str) -> float:
    text = Path(logfile).read_text()
    m1 = re.search(rf"\[(\d{{2}}\.\d{{2}}\|\d{{2}}:\d{{2}}:\d{{2}})\] JOB {jobname} STARTED", text)
    m2 = re.search(rf"\[(\d{{2}}\.\d{{2}}\|\d{{2}}:\d{{2}}:\d{{2}})\] JOB {jobname} FINISHED", text)
    if not (m1 and m2):
        return float("nan")
    fmt = "%d.%m|%H:%M:%S"
    t1 = datetime.strptime(m1.group(1), fmt)
    t2 = datetime.strptime(m2.group(1), fmt)
    return (t2 - t1).total_seconds()


def spectrum_plot(freq: np.ndarray, inten: np.ndarray, title: str, path: str, xlim=(0, 4000)) -> None:
    fig, ax = plt.subplots(figsize=(7, 3))
    ax.vlines(freq, 0, inten, color="C0", linewidth=1.2)
    ax.set_xlim(*xlim)
    ax.set_ylim(0, max(float(np.max(inten)) * 1.1, 1.0))
    ax.set_xlabel("Frequency / cm$^{-1}$")
    ax.set_ylabel("IR intensity / km mol$^{-1}$")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(path, dpi=180)
    plt.close(fig)


def main() -> None:
    init(folder="report_workdir")
    jobs = load_jobs()
    opt, ph, mt, full = jobs["opt"], jobs["partial"], jobs["mt"], jobs["full"]

    f_ph, i_ph, m_ph = vib_arrays(ph, "adf")
    f_mt, i_mt, m_mt = vib_arrays(mt, "ams")
    f_full, i_full, m_full = vib_arrays(full, "adf")

    mt_mode = m_mt[0]
    overlaps = np.array([abs(np.vdot(normalize(mt_mode).ravel(), normalize(m).ravel())) for m in m_full])
    full_match_idx = int(np.argmax(overlaps))

    # Images: selected mode displaced structures, with CO region highlighted.
    system = opt.results.get_main_system()
    system.add_atom_to_region(CIDX, "CO")
    system.add_atom_to_region(OIDX, "CO")
    disp = normalize(mt_mode) * 0.45
    neg = system.copy(); neg.coords = np.array(system.coords) - disp
    pos = system.copy(); pos.coords = np.array(system.coords) + disp
    images = {
        "mode -": view(neg, direction="along_pca3", show_regions=True, width=450, height=350),
        "optimized": view(system, direction="along_pca3", show_regions=True, width=450, height=350),
        "mode +": view(pos, direction="along_pca3", show_regions=True, width=450, height=350),
    }
    plot_image_grid(images, rows=1, save_path="mode_displacement_grid.png")

    spectrum_plot(f_ph, i_ph, "Partial-Hessian spectrum (CO region)", "partial_hessian_spectrum.png")
    spectrum_plot(f_mt, i_mt, "Mode-tracking spectrum (selected C=O mode)", "mode_tracking_spectrum.png")
    spectrum_plot(f_full, i_full, "Full normal-modes IR spectrum", "full_spectrum.png")

    selected = pd.DataFrame([
        {"calculation": "Partial Hessian guess", "mode": PARTIAL_MODE, "frequency_cm-1": f_ph[PARTIAL_MODE-1], "intensity_km/mol": i_ph[PARTIAL_MODE-1], "CO_displacement_fraction": co_fraction(m_ph[PARTIAL_MODE-1])},
        {"calculation": "Mode tracking", "mode": 1, "frequency_cm-1": f_mt[0], "intensity_km/mol": i_mt[0], "CO_displacement_fraction": co_fraction(m_mt[0])},
        {"calculation": "Full normal modes (best overlap)", "mode": full_match_idx + 1, "frequency_cm-1": f_full[full_match_idx], "intensity_km/mol": i_full[full_match_idx], "CO_displacement_fraction": co_fraction(m_full[full_match_idx])},
    ])

    t_opt = parse_duration("01-run_workdir/logfile", "01_opt")
    t_ph = parse_duration("02-continue_workdir/logfile", "02_partial_hessian")
    t_mt = parse_duration("03-rest_workdir/logfile", "03_mode_tracking")
    t_full = parse_duration("03-rest_workdir/logfile", "04_full_normal_modes")
    timings = pd.DataFrame([
        {"job": "01_opt", "description": "Geometry optimization", "wall_s": t_opt},
        {"job": "02_partial_hessian", "description": "CO-region partial Hessian", "wall_s": t_ph},
        {"job": "03_mode_tracking", "description": "Mode tracking from mode 153", "wall_s": t_mt},
        {"job": "04_full_normal_modes", "description": "Full normal modes", "wall_s": t_full},
        {"job": "partial+mode_tracking", "description": "Selective workflow", "wall_s": t_ph + t_mt},
    ])
    speedup = t_full / (t_ph + t_mt)

    md = []
    md.append("# Mode tracking vs full normal modes for dydrogesterone\n")
    md.append("Method: ADF/PBE/DZ with large frozen core. Dydrogesterone was constructed from the requested SMILES. ChemicalSystem bond orders and ring membership identified the non-ring C=O as C2–O3 (1-based), assigned to region `CO`.\n")
    md.append("## Mode displacement image\n\n![mode displacement](mode_displacement_grid.png)\n")
    md.append("## IR spectra\n\n![partial](partial_hessian_spectrum.png)\n\n![mode tracking](mode_tracking_spectrum.png)\n\n![full](full_spectrum.png)\n")
    md.append("## Selected C=O mode\n\n" + selected.to_markdown(index=False, floatfmt=".3f") + "\n")
    md.append("## Timings\n\n" + timings.to_markdown(index=False, floatfmt=".1f") + "\n")
    md.append(f"\nSelective partial-Hessian + mode-tracking time = {t_ph + t_mt:.1f} s vs full normal modes = {t_full:.1f} s, speedup = {speedup:.2f}x.\n")
    md.append("\n## Conclusion\n")
    md.append("Mode tracking gives a C=O frequency and intensity very close to the corresponding full normal-mode result, while requiring much less wall time than the full Hessian.\n")
    md.append("\n## Provenance and inputs\n")
    for key, (path, why) in JOBS.items():
        md.append(f"\n### {key}: `{path}`\n\n{why}\n\n```ams\n{jobs[key].get_input()}\n```\n")
    Path("report.md").write_text("\n".join(md))
    print(Path("report.md").resolve())
    finish()

if __name__ == "__main__":
    main()