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

from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import pandas as pd
from scm.base import Units
from scm.plams import AMSJob, view


WORKDIR = Path("01-benzoyl-peroxide-oo-scan_workdir")
REPORT_PATH = Path("report.md")
PLOT_PATH = Path("pes_scan.png")
BONDED_IMAGE_PATH = Path("bonded_state.png")
DISSOCIATED_IMAGE_PATH = Path("dissociated_state.png")
BOND_CUTOFF_ANGSTROM = 2.0


JOB_METADATA = {
    0: {
        "path": WORKDIR / "spinpol_0" / "ams.rkf",
        "why": "Closed-shell reference PES scan for the peroxide O-O cleavage.",
    },
    1: {
        "path": WORKDIR / "spinpol_1" / "ams.rkf",
        "why": "Open-shell PES scan with one excess alpha electron.",
    },
    2: {
        "path": WORKDIR / "spinpol_2" / "ams.rkf",
        "why": "Open-shell PES scan with two excess alpha electrons.",
    },
}


def load_jobs() -> dict[int, AMSJob]:
    jobs: dict[int, AMSJob] = {}
    for spinpol, metadata in JOB_METADATA.items():
        job = AMSJob.load_external(str(metadata["path"]))
        print(f"Loaded spinpol={spinpol}: {metadata['why']}")
        jobs[spinpol] = job
    return jobs


def build_dataframe(jobs: dict[int, AMSJob]) -> pd.DataFrame:
    rows: list[dict[str, Any]] = []
    for spinpol, job in jobs.items():
        results = job.results.get_pesscan_results(molecules=True)
        conversion = Units.conversion_factor("bohr", "angstrom")
        distances = [distance * conversion for distance in results["RaveledPESCoords"][0]]
        energies = list(results["PES"])
        for point_index, (distance, energy, converged, molecule) in enumerate(
            zip(distances, energies, results["Converged"], results["Molecules"])
        ):
            rows.append(
                {
                    "spinpolarization": spinpol,
                    "point_index": point_index,
                    "distance_angstrom": float(distance),
                    "energy_hartree": float(energy),
                    "converged": bool(converged),
                    "molecule": molecule,
                }
            )
    df = pd.DataFrame(rows)
    global_min = float(df["energy_hartree"].min())
    df["relative_energy_kcal_per_mol"] = (
        df["energy_hartree"] - global_min
    ) * Units.conversion_factor("hartree", "kcal/mol")
    return df


def plot_curves(df: pd.DataFrame) -> None:
    fig, ax = plt.subplots(figsize=(6, 4))
    for spinpol, subdf in df.groupby("spinpolarization"):
        ordered = subdf.sort_values("distance_angstrom")
        ax.plot(
            ordered["distance_angstrom"],
            ordered["relative_energy_kcal_per_mol"],
            marker="o",
            label=f"spinpolarization = {spinpol}",
        )
    ax.set_xlabel("O-O distance [angstrom]")
    ax.set_ylabel("Relative energy [kcal/mol]")
    ax.set_title("Benzoyl peroxide O-O bond PES scan with UMA-S-1.2-OMol")
    ax.legend()
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(PLOT_PATH, dpi=200)
    plt.close(fig)


def select_state_rows(df: pd.DataFrame) -> tuple[pd.Series, pd.Series]:
    bonded_candidates = df[df["distance_angstrom"] <= BOND_CUTOFF_ANGSTROM]
    if bonded_candidates.empty:
        bonded_candidates = df
    bonded_row = bonded_candidates.loc[bonded_candidates["energy_hartree"].idxmin()]
    max_distance = float(df["distance_angstrom"].max())
    dissociated_candidates = df[df["distance_angstrom"] >= max_distance - 1.0e-6]
    dissociated_row = dissociated_candidates.loc[dissociated_candidates["energy_hartree"].idxmin()]
    return bonded_row, dissociated_row


def render_state_images(bonded_row: pd.Series, dissociated_row: pd.Series) -> None:
    view(
        bonded_row["molecule"],
        guess_bonds=True,
        width=350,
        height=260,
        direction="along_pca3",
        picture_path=str(BONDED_IMAGE_PATH),
    )
    view(
        dissociated_row["molecule"],
        guess_bonds=True,
        width=350,
        height=260,
        direction="along_pca3",
        picture_path=str(DISSOCIATED_IMAGE_PATH),
    )


def build_summary_table(df: pd.DataFrame) -> pd.DataFrame:
    summary = (
        df.sort_values(["spinpolarization", "distance_angstrom"])
        .groupby("spinpolarization", as_index=False)
        .agg(
            min_distance_angstrom=("distance_angstrom", "min"),
            max_distance_angstrom=("distance_angstrom", "max"),
            minimum_energy_hartree=("energy_hartree", "min"),
            minimum_relative_energy_kcal_per_mol=("relative_energy_kcal_per_mol", "min"),
        )
    )
    return summary


def write_report(
    df: pd.DataFrame,
    summary: pd.DataFrame,
    bonded_row: pd.Series,
    dissociated_row: pd.Series,
    jobs: dict[int, AMSJob],
) -> None:
    if REPORT_PATH.exists():
        REPORT_PATH.replace(REPORT_PATH.with_suffix(".md.bk"))

    intro = (
        "# Benzoyl Peroxide O-O Bond Scan\n\n"
        "This report summarizes three AMS PES scans for benzoyl peroxide "
        "using the MLPotential engine with model `UMA-S-1.2-OMol`. "
        "The peroxide O-O bond was identified by looping over oxygen atoms "
        "and inspecting `ChemicalSystem.bonds` for an oxygen neighbor.\n\n"
    )

    provenance_lines = []
    for spinpol, metadata in JOB_METADATA.items():
        provenance_lines.append(
            f"- `spinpol_{spinpol}` loaded from `{metadata['path']}`. {metadata['why']}"
        )

    summary_table = summary.to_markdown(index=False, floatfmt=".6f")

    lowest_states = (
        "## Lowest-Energy Structures\n\n"
        f"The lowest-energy bonded state in the full scan set occurs at "
        f"`spinpolarization = {int(bonded_row['spinpolarization'])}` and "
        f"`d(O-O) = {bonded_row['distance_angstrom']:.3f}` angstrom. "
        f"Here, 'bonded' means `d(O-O) <= {BOND_CUTOFF_ANGSTROM:.1f}` angstrom.\n\n"
        f"![Lowest-energy bonded state]({BONDED_IMAGE_PATH})\n\n"
        f"The lowest-energy dissociated state was defined as the lowest-energy "
        f"structure at the largest scanned O-O distance "
        f"(`d(O-O) = {dissociated_row['distance_angstrom']:.3f}` angstrom). "
        f"It occurs for `spinpolarization = {int(dissociated_row['spinpolarization'])}`.\n\n"
        f"![Lowest-energy dissociated state]({DISSOCIATED_IMAGE_PATH})\n\n"
    )

    conclusion = (
        "## Conclusion\n\n"
        f"Among the three scans, the global minimum is found for "
        f"`spinpolarization = {int(bonded_row['spinpolarization'])}`. "
        f"At the dissociation limit sampled here (`~3.0` angstrom), the lowest-energy "
        f"state is obtained for `spinpolarization = {int(dissociated_row['spinpolarization'])}`.\n"
    )

    inputs = []
    for spinpol, job in jobs.items():
        inputs.append(f"### spinpol_{spinpol}\n\n```text\n{job.get_input()}\n```\n")

    report_text = (
        intro
        + "## Provenance\n\n"
        + "\n".join(provenance_lines)
        + "\n\n## Energy Plot\n\n"
        + f"![Energy vs O-O distance]({PLOT_PATH})\n\n"
        + "## Summary Table\n\n"
        + summary_table
        + "\n\n"
        + lowest_states
        + "## Calculation Inputs\n\n"
        + "\n".join(inputs)
        + conclusion
    )
    REPORT_PATH.write_text(report_text)


def main() -> None:
    jobs = load_jobs()
    df = build_dataframe(jobs)
    summary = build_summary_table(df)
    plot_curves(df)
    bonded_row, dissociated_row = select_state_rows(df)
    render_state_images(bonded_row, dissociated_row)
    write_report(df, summary, bonded_row, dissociated_row, jobs)


if __name__ == "__main__":
    main()