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

import re
from dataclasses import dataclass
from pathlib import Path

import pandas as pd
from scm.base import ChemicalSystem, Units
from scm.conformers import ConformersJob
from scm.plams import Settings, view


ROOT = Path(__file__).resolve().parent


@dataclass(frozen=True)
class Isomer:
    name: str
    kind: str
    smiles: str
    slug: str


ISOMERS = (
    Isomer("butan-1-ol", "alcohol", "CCCCO", "butan_1_ol"),
    Isomer("butan-2-ol", "alcohol", "CCC(O)C", "butan_2_ol"),
    Isomer("2-methylpropan-1-ol", "alcohol", "CC(C)CO", "2_methylpropan_1_ol"),
    Isomer("2-methylpropan-2-ol", "alcohol", "CC(C)(C)O", "2_methylpropan_2_ol"),
    Isomer("ethoxyethane (diethyl ether)", "ether", "CCOCC", "ethoxyethane"),
    Isomer(
        "1-methoxypropane (methyl n-propyl ether)",
        "ether",
        "COCCC",
        "1_methoxypropane",
    ),
    Isomer(
        "2-methoxypropane (methyl isopropyl ether)",
        "ether",
        "COC(C)C",
        "2_methoxypropane",
    ),
)


def workdir_generation(path: Path) -> int:
    match = re.fullmatch(r"01-conformers_workdir(?:\.(\d+))?", path.name)
    if match is None:
        return -1
    return int(match.group(1) or 1)


def latest_workdir() -> Path:
    candidates = [
        path
        for path in ROOT.glob("01-conformers_workdir*")
        if path.is_dir() and workdir_generation(path) >= 0
    ]
    if not candidates:
        raise FileNotFoundError("No 01-conformers_workdir job directory found")
    return max(candidates, key=workdir_generation)


def conformer_settings() -> Settings:
    settings = Settings()
    settings.input.ams.Task = "Generate"
    settings.input.ams.Generator.Method = "RDKit"
    settings.input.ams.Generator.RDKit.InitialNConformers = 8
    settings.input.MLPotential.Model = "AIMNet2-wB97MD3"
    settings.runscript.nproc = 1
    settings.runscript.pre = "export OMP_NUM_THREADS=1"
    return settings


def load_jobs(workdir: Path) -> dict[str, ConformersJob]:
    jobs: dict[str, ConformersJob] = {}
    for isomer in ISOMERS:
        rkf = workdir / isomer.slug / "conformers.rkf"
        if not rkf.is_file():
            raise FileNotFoundError(f"Missing conformer result: {rkf}")
        jobs[isomer.slug] = ConformersJob.load_external(
            str(rkf),
            settings=conformer_settings(),
            molecule=ChemicalSystem.from_smiles(isomer.smiles),
        )
    return jobs


def render_report() -> None:
    workdir = latest_workdir()
    jobs = load_jobs(workdir)
    tables_dir = ROOT / "tables"
    figures_dir = ROOT / "figures"
    tables_dir.mkdir(exist_ok=True)
    figures_dir.mkdir(exist_ok=True)

    hartree_to_kcal_mol = Units.conversion_factor("hartree", "kcal/mol")
    rows: list[dict[str, object]] = []
    for isomer in ISOMERS:
        job = jobs[isomer.slug]
        conformers = job.results.get_conformers()
        lowest = conformers[0]
        picture_path = figures_dir / f"{isomer.slug}.png"
        view(
            lowest,
            guess_bonds=len(lowest.bonds) == 0,
            direction="along_pca3",
            width=360,
            height=280,
            picture_path=str(picture_path),
        )
        rows.append(
            {
                "Name": isomer.name,
                "Type": isomer.kind,
                "SMILES": isomer.smiles,
                "Picture": f'<img src="figures/{picture_path.name}" width="180">',
                "Electronic energy (hartree)": job.results.get_lowest_energy("au"),
                "Unique conformers": len(conformers),
            }
        )

    frame = pd.DataFrame(rows)
    minimum_hartree = float(frame["Electronic energy (hartree)"].min())
    frame["Relative energy (kcal/mol)"] = (
        frame["Electronic energy (hartree)"] - minimum_hartree
    ) * hartree_to_kcal_mol
    frame = frame.sort_values("Relative energy (kcal/mol)", kind="stable").reset_index(drop=True)

    csv_frame = frame.copy()
    csv_frame["Picture"] = csv_frame["Picture"].str.extract(r'src="([^"]+)"', expand=False)
    csv_frame.to_csv(tables_dir / "isomer_relative_energies.csv", index=False)

    display = frame[["Name", "SMILES", "Picture", "Relative energy (kcal/mol)"]].copy()
    display["Relative energy (kcal/mol)"] = display["Relative energy (kcal/mol)"].map(
        lambda value: f"{value:.3f}"
    )
    winner = str(frame.iloc[0]["Name"])

    lines = [
        "# C4H10O constitutional-isomer stability",
        "",
        "Seven neutral, closed-shell gas-phase constitutional isomers were compared with "
        "AIMNet2-wB97MD3. For each isomer, RDKit ETKDG supplied eight starting geometries. "
        "The AMS Conformers tool optimized and ranked them with the same AIMNet2 model. "
        "Duplicate minima were removed by the Conformers defaults.",
        "",
        "The table uses the electronic energy of the lowest-energy optimized conformer for "
        "each isomer. It contains no zero-point, thermal, enthalpic, entropic, or other "
        "vibrational correction. Relative energies share the minimum electronic energy of "
        "the seven isomers as their zero.",
        "",
        display.to_markdown(index=False, disable_numparse=True),
        "",
        f"The lowest electronic energy in this calculation belongs to **{winner}**.",
        "",
        "## Calculation details and provenance",
        "",
        f"Results were loaded directly from `{workdir.name}`. AIMNet2-wB97MD3 is trained "
        "against ωB97M-D3/def2-TZVPP reference data. The calculations used one process and "
        "one OpenMP thread per job, and the seven jobs ran sequentially.",
        "",
        "The `Unique conformers` column in the CSV records how many distinct optimized "
        "structures remained from the eight initial geometries.",
        "",
    ]
    for isomer in ISOMERS:
        job = jobs[isomer.slug]
        lines.extend(
            [
                f"### {isomer.name}",
                "",
                f"Loaded because it provides the optimized conformers and electronic "
                f"energies for `{isomer.smiles}`.",
                "",
                "```ams",
                job.get_input().rstrip(),
                "```",
                "",
            ]
        )

    (ROOT / "report.md").write_text("\n".join(lines), encoding="utf-8")


if __name__ == "__main__":
    render_report()