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

from dataclasses import dataclass
from pathlib import Path
import shutil

import pandas as pd
import scm.plams as plams
from scm.base import Units
from scm.conformers import ConformersJob


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


MOLECULES: tuple[MoleculeSpec, ...] = (
    MoleculeSpec("butan-1-ol", "CCCCO", "butan-1-ol", "alcohol"),
    MoleculeSpec("butan-2-ol", "CCC(O)C", "butan-2-ol", "alcohol"),
    MoleculeSpec("2-methylpropan-1-ol", "CC(C)CO", "2-methylpropan-1-ol", "alcohol"),
    MoleculeSpec("2-methylpropan-2-ol", "CC(C)(C)O", "2-methylpropan-2-ol", "alcohol"),
    MoleculeSpec("ethoxyethane (diethyl ether)", "CCOCC", "ethoxyethane", "ether"),
    MoleculeSpec("1-methoxypropane (methyl n-propyl ether)", "COCCC", "1-methoxypropane", "ether"),
    MoleculeSpec("2-methoxypropane (methyl isopropyl ether)", "COC(C)C", "2-methoxypropane", "ether"),
)


def load_rows(images_dir: Path) -> tuple[pd.DataFrame, list[str]]:
    rows: list[dict[str, object]] = []
    provenance: list[str] = []
    hartree_to_kcal = Units.conversion_factor("Ha", "kcal/mol")
    job_root = Path("01-run_workdir")

    for spec in MOLECULES:
        job_path = job_root / spec.slug
        job = ConformersJob.load_external(job_path)
        provenance.append(
            f"- `{spec.slug}` loaded from `{job_path}`: 8 RDKit-started conformers evaluated with `MLPotential/AIMNet2-wB97MD3`; lowest-energy conformer used for ranking."
        )
        lowest = job.results.get_lowest_conformer()
        image_name = f"{spec.slug}.png"
        image_path = images_dir / image_name
        plams.view(lowest, width=220, height=180, guess_bonds=True, direction="along_pca3", picture_path=str(image_path))
        rows.append(
            {
                "name": spec.name,
                "type": spec.kind,
                "smiles": spec.smiles,
                "picture": f"![{spec.name}]({image_path.as_posix()})",
                "energy_ha": job.results.get_lowest_energy("Ha"),
                "input": job.get_input().strip(),
            }
        )

    df = pd.DataFrame(rows)
    min_energy = float(df["energy_ha"].min())
    df["relative_energy_kcal_mol"] = (df["energy_ha"] - min_energy) * hartree_to_kcal
    df.sort_values(["relative_energy_kcal_mol", "name"], inplace=True, ignore_index=True)
    return df, provenance


def dataframe_to_markdown(df: pd.DataFrame) -> str:
    columns = ["name", "smiles", "picture", "relative_energy_kcal_mol"]
    md = df.loc[:, columns].rename(
        columns={
            "name": "Name",
            "smiles": "SMILES",
            "picture": "Picture",
            "relative_energy_kcal_mol": "Relative Energy (kcal/mol)",
        }
    ).copy()
    md["Relative Energy (kcal/mol)"] = md["Relative Energy (kcal/mol)"].map(lambda x: f"{x:.3f}")
    return md.to_markdown(index=False)


def write_report(df: pd.DataFrame, provenance: list[str], report_path: Path) -> None:
    if report_path.exists():
        shutil.copy2(report_path, report_path.with_suffix(".md.bk"))

    lowest_name = str(df.iloc[0]["name"])
    intro = (
        "# C4H10O Constitutional Isomer Ranking\n\n"
        "This report ranks the seven listed constitutional isomers of `C4H10O` by the electronic energy of their lowest-energy conformer.\n"
        "For each molecule, AMS Conformers was run with `Generator=RDKit`, `InitialNConformers=8`, and `MLPotential/AIMNet2-wB97MD3`.\n"
        "Vibrational effects were excluded, as requested.\n"
    )
    methods = "## Provenance\n\n" + "\n".join(provenance) + "\n"
    table = "## Results\n\n" + dataframe_to_markdown(df) + "\n"
    inputs = "\n".join(
        f"### {row['name']}\n\n```text\n{row['input']}\n```\n" for _, row in df.iterrows()
    )
    conclusion = (
        "## Conclusion\n\n"
        f"The lowest-energy structure in this set is **{lowest_name}**, so it is the predicted most stable constitutional isomer at the `AIMNet2-wB97MD3` level used here.\n"
    )
    report_path.write_text("\n\n".join([intro, methods, table, "## Calculation Inputs\n\n" + inputs, conclusion]))


def main() -> None:
    plams.init()
    images_dir = Path("images")
    images_dir.mkdir(exist_ok=True)
    df, provenance = load_rows(images_dir)
    write_report(df, provenance, Path("report.md"))
    print(df.loc[:, ["name", "relative_energy_kcal_mol"]].to_string(index=False))
    plams.finish()


if __name__ == "__main__":
    main()