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

from dataclasses import dataclass
import os
from pathlib import Path
import shutil
import subprocess

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scm.plams import AMSJob


@dataclass(frozen=True)
class JobRef:
    label: str
    method: str
    path: Path
    justification: str


def latest_workdir(prefix: str) -> Path:
    candidates = [p for p in Path(".").glob(f"{prefix}*") if p.is_dir()]
    if not candidates:
        raise FileNotFoundError(f"No PLAMS workdir matching {prefix}*")
    return max(candidates, key=lambda p: p.stat().st_mtime)


def frontier_window(job: AMSJob, n_each_side: int = 3) -> tuple[pd.DataFrame, dict[str, float]]:
    energies = np.asarray(job.results.get_orbital_energies(unit="eV"))[0]
    occupations = np.asarray(job.results.get_orbital_occupations())[0]
    occupied = np.where(occupations > 1.0e-6)[0]
    virtual = np.where(occupations <= 1.0e-6)[0]
    if len(occupied) == 0 or len(virtual) == 0:
        raise RuntimeError(f"Could not identify occupied and virtual orbitals for {job.name}")

    homo_idx = int(occupied[-1])
    lumo_idx = int(virtual[virtual > homo_idx][0])
    start = max(0, homo_idx - n_each_side)
    stop = min(len(energies), lumo_idx + n_each_side + 1)

    rows: list[dict[str, object]] = []
    for idx in range(start, stop):
        label = f"HOMO{idx - homo_idx:+d}" if idx <= homo_idx else f"LUMO{idx - lumo_idx:+d}"
        label = label.replace("+0", "").replace("-0", "")
        rows.append(
            {
                "Orbital": label,
                "Index": idx + 1,
                "Occupation": float(occupations[idx]),
                "Energy (eV)": float(energies[idx]),
            }
        )

    summary = {
        "HOMO (eV)": float(job.results.get_homo_energies(unit="eV")[0]),
        "LUMO (eV)": float(job.results.get_lumo_energies(unit="eV")[0]),
        "Gap (eV)": float(job.results.get_smallest_homo_lumo_gap(unit="eV")),
        "HOMO index": homo_idx + 1,
        "LUMO index": lumo_idx + 1,
    }
    return pd.DataFrame(rows), summary


def plot_levels(windows: dict[str, pd.DataFrame], output: Path) -> None:
    fig, ax = plt.subplots(figsize=(5.8, 3.6))
    x_positions = {"PBE/DZP": 0.0, "B3LYP/DZP": 1.0}
    colors = {"occupied": "#1f77b4", "virtual": "#d62728"}

    for method, df in windows.items():
        x = x_positions[method]
        for _, row in df.iterrows():
            occ = float(row["Occupation"])
            energy = float(row["Energy (eV)"])
            color = colors["occupied" if occ > 1.0e-6 else "virtual"]
            ax.hlines(energy, x - 0.22, x + 0.22, color=color, linewidth=2)
            if row["Orbital"] in {"HOMO", "LUMO"}:
                ax.text(x + 0.26, energy, str(row["Orbital"]), va="center", fontsize=8)

    ax.set_xticks(list(x_positions.values()), list(x_positions.keys()))
    ax.set_ylabel("Orbital energy (eV)")
    ax.set_xlim(-0.55, 1.65)
    ax.grid(axis="y", color="#dddddd", linewidth=0.8)
    ax.spines[["top", "right"]].set_visible(False)
    ax.plot([], [], color=colors["occupied"], label="Occupied")
    ax.plot([], [], color=colors["virtual"], label="Virtual")
    ax.legend(frameon=False, fontsize=8, loc="lower right")
    fig.tight_layout()
    fig.savefig(output, dpi=180)
    plt.close(fig)


def generate_amsreport_image(result_file: Path, orbital: str, html_file: Path, image_file: Path) -> None:
    amsbin = Path(os.environ["AMSBIN"])
    command = [
        "xvfb-run",
        "-a",
        str(amsbin / "amsreport"),
        str(result_file),
        orbital,
        "-o",
        str(html_file),
        "-v",
        "-grid Fine",
        "-v",
        "-antialias",
        "-v",
        "-bgcolor #ffffff",
        "-v",
        "-scmgeometry 500x500",
    ]
    subprocess.run(command, check=True)
    generated = html_file.with_suffix(".jpgs") / "0.jpg"
    if not generated.exists() or generated.stat().st_size == 0:
        raise RuntimeError(f"AMSreport did not create a usable image for {orbital} from {result_file}")
    shutil.copyfile(generated, image_file)


def fmt_df(df: pd.DataFrame, digits: int = 3) -> str:
    return df.to_markdown(index=False, floatfmt=f".{digits}f")


def main() -> None:
    workdir = latest_workdir("01-run_workdir")
    refs = [
        JobRef(
            "Geometry optimization",
            "ADF/PBE/DZP",
            workdir / "ethylene_opt_pbe_dzp",
            "Optimized the ethylene geometry once from SMILES C=C.",
        ),
        JobRef(
            "GGA single point",
            "ADF/PBE/DZP",
            workdir / "ethylene_sp_pbe_dzp",
            "Computed frontier orbital energies on the optimized geometry with the GGA setup.",
        ),
        JobRef(
            "Hybrid single point",
            "ADF/B3LYP/DZP",
            workdir / "ethylene_sp_b3lyp_dzp",
            "Computed frontier orbital energies on the same optimized geometry with a hybrid functional.",
        ),
    ]
    jobs = {ref.label: AMSJob.load_external(str(ref.path)) for ref in refs}

    windows: dict[str, pd.DataFrame] = {}
    summary_rows: list[dict[str, object]] = []
    for ref in refs[1:]:
        df, summary = frontier_window(jobs[ref.label])
        method_key = "PBE/DZP" if "PBE" in ref.method else "B3LYP/DZP"
        windows[method_key] = df
        summary_rows.append({"Method": ref.method, **summary})

    figures = Path("figures")
    figures.mkdir(exist_ok=True)
    plot_levels(windows, figures / "frontier_energy_levels.png")

    amsreport_dir = Path("amsreport_images")
    amsreport_dir.mkdir(exist_ok=True)
    generate_amsreport_image(
        refs[1].path / "adf.rkf",
        "HOMO",
        amsreport_dir / "pbe_homo.html",
        figures / "pbe_homo_amsreport.jpg",
    )
    generate_amsreport_image(
        refs[1].path / "adf.rkf",
        "LUMO",
        amsreport_dir / "pbe_lumo.html",
        figures / "pbe_lumo_amsreport.jpg",
    )
    generate_amsreport_image(
        refs[2].path / "adf.rkf",
        "HOMO",
        amsreport_dir / "b3lyp_homo.html",
        figures / "b3lyp_homo_amsreport.jpg",
    )
    generate_amsreport_image(
        refs[2].path / "adf.rkf",
        "LUMO",
        amsreport_dir / "b3lyp_lumo.html",
        figures / "b3lyp_lumo_amsreport.jpg",
    )

    opt_system = jobs["Geometry optimization"].results.get_main_system()
    cc_distance = opt_system.get_distance(0, 1, unit="angstrom")
    ch_distances = [opt_system.get_distance(0, i, unit="angstrom") for i in (2, 3)]
    ch_distances += [opt_system.get_distance(1, i, unit="angstrom") for i in (4, 5)]

    summary_df = pd.DataFrame(summary_rows)
    report = [
        "# Ethylene ADF Frontier Orbital Comparison",
        "",
        "Ethylene was built from SMILES `C=C`. The geometry was optimized once with a modest ADF GGA setup, then two ADF single-point calculations were run on the optimized geometry: PBE/DZP and B3LYP/DZP. Orbital energies are reported in eV.",
        "",
        "## Job provenance",
        "",
        "| Role | Method | Job directory | Justification |",
        "|---|---|---|---|",
    ]
    for ref in refs:
        report.append(f"| {ref.label} | {ref.method} | `{ref.path}` | {ref.justification} |")

    report += [
        "",
        "## Optimized geometry",
        "",
        f"The optimized C=C distance is {cc_distance:.3f} angstrom. The mean C-H distance is {np.mean(ch_distances):.3f} angstrom.",
        "",
        "![Optimized ethylene geometry](figures/optimized_ethylene.png)",
        "",
        "## HOMO-LUMO comparison",
        "",
        fmt_df(summary_df[["Method", "HOMO (eV)", "LUMO (eV)", "Gap (eV)", "HOMO index", "LUMO index"]]),
        "",
        "## Frontier orbital energy windows",
        "",
        "### ADF/PBE/DZP",
        "",
        fmt_df(windows["PBE/DZP"]),
        "",
        "### ADF/B3LYP/DZP",
        "",
        fmt_df(windows["B3LYP/DZP"]),
        "",
        "## Energy-level diagram",
        "",
        "![Frontier orbital energy levels](figures/frontier_energy_levels.png)",
        "",
        "## AMSreport orbital isosurfaces",
        "",
        "These images were generated by `amsreport` from each single-point `adf.rkf` file.",
        "",
        "| Method | HOMO | LUMO |",
        "|---|---|---|",
        "| ADF/PBE/DZP | ![PBE HOMO](figures/pbe_homo_amsreport.jpg) | ![PBE LUMO](figures/pbe_lumo_amsreport.jpg) |",
        "| ADF/B3LYP/DZP | ![B3LYP HOMO](figures/b3lyp_homo_amsreport.jpg) | ![B3LYP LUMO](figures/b3lyp_lumo_amsreport.jpg) |",
        "",
        "## Calculation inputs",
        "",
    ]
    for ref in refs:
        report += [
            f"### {ref.label}: {ref.method}",
            "",
            "```ams",
            jobs[ref.label].get_input(),
            "```",
            "",
        ]

    report += [
        "## Conclusion",
        "",
        f"With this modest DZP setup, the B3LYP single point increases the HOMO-LUMO gap from {summary_rows[0]['Gap (eV)']:.3f} eV for PBE to {summary_rows[1]['Gap (eV)']:.3f} eV. The increase is mainly from a lower HOMO and a higher LUMO relative to the PBE values.",
        "",
    ]
    Path("report.md").write_text("\n".join(report), encoding="utf-8")
    print(Path("report.md").resolve())


if __name__ == "__main__":
    main()