#!/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
TABLES = ROOT / "tables"
FIGURES = ROOT / "figures"
JOB_NAME = "formaldehyde_adf_pbe_dzp"


def latest_job_path() -> Path:
    workdirs = [path for path in ROOT.glob("01-run_workdir*") if path.is_dir()]
    if not workdirs:
        raise FileNotFoundError("No 01-run_workdir directory was found")
    latest = max(workdirs, key=lambda path: path.stat().st_mtime)
    job_path = latest / JOB_NAME
    if not job_path.is_dir():
        raise FileNotFoundError(f"Expected AMS job directory: {job_path}")
    return job_path


def atom_indices(system) -> tuple[int, int, list[int]]:
    carbon = [i for i, atom in enumerate(system.atoms) if atom.symbol == "C"]
    oxygen = [i for i, atom in enumerate(system.atoms) if atom.symbol == "O"]
    hydrogens = [i for i, atom in enumerate(system.atoms) if atom.symbol == "H"]
    if len(carbon) != 1 or len(oxygen) != 1 or len(hydrogens) != 2:
        raise ValueError("The optimized system does not have the expected CH2O composition")
    return carbon[0], oxygen[0], hydrogens


def symmetry_labels(job: AMSJob, count: int) -> list[str]:
    raw = job.results.readrkf("Vibrations", "IrReps", file="engine")
    if isinstance(raw, str):
        labels = raw.split()
    else:
        labels = [str(value).strip() for value in np.asarray(raw).reshape(-1)]
    if len(labels) != count:
        raise ValueError(f"Found {len(labels)} symmetry labels for {count} modes")
    return labels


def internal_derivatives(system, mode: np.ndarray) -> dict[str, float]:
    c_idx, o_idx, h_idx = atom_indices(system)
    step = 0.01
    plus = system.copy()
    minus = system.copy()
    plus.coords[:] = np.asarray(system.coords) + step * mode
    minus.coords[:] = np.asarray(system.coords) - step * mode

    def central(value_plus: float, value_minus: float) -> float:
        return (value_plus - value_minus) / (2.0 * step)

    d_co = central(plus.get_distance(c_idx, o_idx), minus.get_distance(c_idx, o_idx))
    d_ch = [
        central(plus.get_distance(c_idx, h), minus.get_distance(c_idx, h))
        for h in h_idx
    ]
    d_angle = central(
        plus.get_angle(h_idx[0], c_idx, h_idx[1], unit="degree"),
        minus.get_angle(h_idx[0], c_idx, h_idx[1], unit="degree"),
    )
    coords = np.asarray(system.coords)
    molecular_normal = np.cross(coords[o_idx] - coords[c_idx], coords[h_idx[0]] - coords[c_idx])
    molecular_normal /= np.linalg.norm(molecular_normal)
    oop = sum(abs(np.dot(mode[h] - mode[c_idx], molecular_normal)) for h in h_idx)
    return {"co": d_co, "ch1": d_ch[0], "ch2": d_ch[1], "angle": d_angle, "oop": oop}


def assign_modes(frequencies: np.ndarray, modes: np.ndarray, system) -> list[str]:
    metrics = [internal_derivatives(system, mode) for mode in modes]
    assignments = [""] * len(frequencies)
    high = [i for i, frequency in enumerate(frequencies) if frequency > 2200.0]
    low = [i for i in range(len(frequencies)) if i not in high]

    for i in high:
        same_phase = metrics[i]["ch1"] * metrics[i]["ch2"] >= 0.0
        assignments[i] = "symmetric C-H stretch" if same_phase else "asymmetric C-H stretch"

    if low:
        oop_idx = max(low, key=lambda i: metrics[i]["oop"])
        assignments[oop_idx] = "CH2 out-of-plane bend"
        remaining = [i for i in low if i != oop_idx]
        if remaining:
            co_idx = max(remaining, key=lambda i: abs(metrics[i]["co"]))
            assignments[co_idx] = "C=O stretch"
            remaining.remove(co_idx)
        if remaining:
            scissor_idx = max(remaining, key=lambda i: abs(metrics[i]["angle"]))
            assignments[scissor_idx] = "H-C-H scissoring bend"
            remaining.remove(scissor_idx)
        for i in remaining:
            assignments[i] = "CH2 in-plane rocking bend"

    return [assignment or "mixed vibration" for assignment in assignments]


def make_mode_grids(system, modes: np.ndarray) -> list[Path]:
    grid_paths: list[Path] = []
    optimized_image = FIGURES / "optimized_structure.png"
    view(
        system,
        guess_bonds=len(system.bonds) == 0,
        direction="along_pca3",
        width=500,
        height=400,
        picture_path=str(optimized_image),
    )

    for number, mode in enumerate(modes, start=1):
        max_norm = float(np.linalg.norm(mode, axis=1).max())
        scale = 0.25 / max_norm
        negative = system.copy()
        positive = system.copy()
        negative.coords[:] = np.asarray(system.coords) - scale * mode
        positive.coords[:] = np.asarray(system.coords) + scale * mode

        images = {}
        for label, structure in [
            ("negative displacement", negative),
            ("optimized structure", system),
            ("positive displacement", positive),
        ]:
            image_path = FIGURES / f"mode_{number:02d}_{label.replace(' ', '_')}.png"
            images[label] = view(
                structure,
                guess_bonds=len(structure.bonds) == 0,
                direction="along_pca3",
                width=360,
                height=300,
                picture_path=str(image_path),
            )
        grid_path = FIGURES / f"mode_{number:02d}_grid.png"
        plot_image_grid(images, rows=1, figsize=(12, 3.5), save_path=str(grid_path))
        plt.close("all")
        grid_paths.append(grid_path)
    return grid_paths


def markdown_image(path: Path, alt: str) -> str:
    return f"![{alt}]({path.relative_to(ROOT).as_posix()})"


def format_input(input_text: str) -> str:
    return f"```ams\n{input_text.rstrip()}\n```"


def main() -> None:
    TABLES.mkdir(exist_ok=True)
    FIGURES.mkdir(exist_ok=True)

    job_path = latest_job_path()
    job = AMSJob.load_external(str(job_path))
    provenance = "Combined ADF geometry optimization, harmonic frequencies, IR intensities, and PES-point characterization."

    system = job.results.get_main_system()
    symmetry_probe = system.copy()
    optimized_point_group = symmetry_probe.symmetrize_molecule(tolerance=0.05)
    frequencies = np.asarray(job.results.get_frequencies(unit="cm^-1"), dtype=float)
    intensities = np.asarray(job.results.get_ir_intensities(), dtype=float)
    modes = np.asarray(job.results.get_normal_modes(), dtype=float)
    labels = symmetry_labels(job, len(frequencies))
    assignments = assign_modes(frequencies, modes, system)

    c_idx, o_idx, h_idx = atom_indices(system)
    co_length = system.get_distance(c_idx, o_idx)
    ch_lengths = [system.get_distance(c_idx, h) for h in h_idx]
    hch_angle = system.get_angle(h_idx[0], c_idx, h_idx[1], unit="degree")

    try:
        pes_character = str(job.results.readrkf("AMSResults", "PESPointCharacter", file="engine")).strip()
    except KeyError:
        pes_character = str(job.results.readrkf("AMSResults", "PESPointCharacter", file="ams")).strip()
    imaginary = frequencies[frequencies < 0.0]
    is_minimum = len(imaginary) == 0 and "minimum" in pes_character.lower()

    frequency_table = pd.DataFrame(
        {
            "Mode": np.arange(1, len(frequencies) + 1),
            "Frequency (cm^-1)": np.round(frequencies, 2),
            "IR intensity (km mol^-1)": np.round(intensities, 3),
            "Symmetry": labels,
            "Qualitative assignment": assignments,
        }
    )
    frequency_table.to_csv(TABLES / "vibrational_modes.csv", index=False)

    geometry_table = pd.DataFrame(
        {
            "Quantity": ["C=O", "C-H(1)", "C-H(2)", "mean C-H", "H-C-H"],
            "Value": [co_length, ch_lengths[0], ch_lengths[1], np.mean(ch_lengths), hch_angle],
            "Unit": ["angstrom", "angstrom", "angstrom", "angstrom", "degree"],
        }
    )
    geometry_table.to_csv(TABLES / "optimized_geometry.csv", index=False)

    spectrum_x, spectrum_y = job.results.get_ir_spectrum(
        broadening_type="gaussian",
        broadening_width=20,
        min_x=400,
        max_x=3400,
        x_spacing=1.0,
    )
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(spectrum_x, spectrum_y, color="#214f8b", linewidth=1.6)
    ax.vlines(frequencies, 0.0, intensities, color="#a33a2b", linewidth=0.8, alpha=0.7)
    ax.set_xlim(3400, 400)
    ax.set_xlabel(r"Wavenumber (cm$^{-1}$)")
    ax.set_ylabel(r"IR intensity (km mol$^{-1}$)")
    ax.set_title("Harmonic IR spectrum, Gaussian width 20 cm$^{-1}$")
    fig.tight_layout()
    spectrum_path = FIGURES / "ir_spectrum.png"
    fig.savefig(spectrum_path, dpi=180)
    plt.close(fig)

    mode_grid_paths = make_mode_grids(system, modes)
    mode_sections = []
    for row, grid_path in zip(frequency_table.itertuples(index=False), mode_grid_paths):
        mode_sections.append(
            f"### Mode {row.Mode}: {row._1:.2f} cm^-1, {row.Symmetry}\n\n"
            f"Assignment: {row._4}. The left and right structures use equal displacements along the normal coordinate.\n\n"
            f"{markdown_image(grid_path, f'Mode {row.Mode} displacement grid')}"
        )

    if is_minimum:
        conclusion = (
            "The optimized formaldehyde structure is a local minimum. AMS classified the PES point as "
            f"{pes_character}, and all {len(frequencies)} vibrational frequencies are positive."
        )
    else:
        conclusion = (
            "The minimum test did not pass. The PES classification is "
            f"{pes_character}, and the number of negative frequencies is {len(imaginary)}."
        )

    report = f"""# ADF optimization and vibrational analysis of formaldehyde

## Calculation

Formaldehyde was built from the SMILES string `C=O`, symmetrized, and optimized as a neutral closed-shell molecule. The optimized structure has {optimized_point_group} symmetry. ADF used PBE-D3(BJ), the all-electron DZP basis, and default numerical and SCF settings. Harmonic normal modes and IR intensities were calculated at the final optimized geometry.

Job directory: `{job_path.relative_to(ROOT)}`

Provenance: {provenance}

## Optimized geometry

{geometry_table.to_markdown(index=False, floatfmt='.6f')}

{markdown_image(FIGURES / 'optimized_structure.png', 'Optimized formaldehyde structure')}

## Minimum verification

AMS PES-point classification: `{pes_character}`.

Negative harmonic frequencies: {len(imaginary)}.

{conclusion}

## Vibrational modes

The assignments use changes in C=O distance, the two C-H distances, H-C-H angle, and displacement normal to the molecular plane. They are qualitative harmonic-mode descriptions.

{frequency_table.to_markdown(index=False)}

## Broadened IR spectrum

The blue curve applies Gaussian broadening with a width of 20 cm^-1 to the calculated sticks. Red lines mark the unbroadened modes. The wavenumber axis follows the usual IR convention.

{markdown_image(spectrum_path, 'Broadened harmonic IR spectrum')}

## Normal-mode displacement pictures

The maximum atomic displacement in each outer structure is 0.25 angstrom. All views use `along_pca3`.

{chr(10).join(mode_sections)}

## AMS input

{format_input(job.get_input())}

## Conclusion

{conclusion} The optimized bond metrics are C=O {co_length:.6f} angstrom, mean C-H {np.mean(ch_lengths):.6f} angstrom, and H-C-H {hch_angle:.4f} degrees.
"""
    (ROOT / "report.md").write_text(report, encoding="utf-8")
    print(f"Wrote {ROOT / 'report.md'}")
    print(conclusion)


if __name__ == "__main__":
    main()