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

from dataclasses import dataclass
from pathlib import Path
import shutil

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scm.plams as plams
from scm.base import ChemicalSystem, Units
from scm.plams import AMSJob


SMILES = "C=CC(CCCC)C1C=CC(C)=CC1=O"
WORKDIR = Path("01-run_workdir")
REPORT_PATH = Path("report.md")
BACKUP_PATH = Path("report.md.bk")


@dataclass(frozen=True)
class LoadedJob:
    job: AMSJob
    why: str


def load_jobs() -> dict[str, LoadedJob]:
    specs = {
        "reactant": LoadedJob(
            job=AMSJob.load_external(str(WORKDIR / "reactant_opt")),
            why="Optimized reactant with normal modes for the reference Gibbs free energy.",
        ),
        "scan": LoadedJob(
            job=AMSJob.load_external(str(WORKDIR / "pes_scan")),
            why="Constrained PES scan used to generate the scan profile and TS starting geometry.",
        ),
        "ts": LoadedJob(
            job=AMSJob.load_external(str(WORKDIR / "ts_search")),
            why="Transition-state search with initial Hessian and normal-mode verification.",
        ),
        "product": LoadedJob(
            job=AMSJob.load_external(str(WORKDIR / "product_opt")),
            why="Optimized product candidate taken from the shortest-distance scan structure.",
        ),
    }
    return specs


def energy_hartree(job: AMSJob) -> float:
    return float(job.results.get_energy(unit="hartree"))


def gibbs_hartree(job: AMSJob) -> float:
    return float(job.results.readrkf("Thermodynamics", "Gibbs free Energy", file="engine"))


def relative_kcal(values_hartree: dict[str, float], reference_key: str) -> dict[str, float]:
    factor = Units.conversion_factor("hartree", "kcal/mol")
    reference = values_hartree[reference_key]
    return {key: (value - reference) * factor for key, value in values_hartree.items()}


def render_system_image(system: ChemicalSystem, picture_path: Path) -> None:
    plams.view(
        system,
        guess_bonds=True,
        direction="tilt_pca3",
        width=360,
        height=280,
        picture_path=str(picture_path),
    )


def write_report(markdown: str) -> None:
    if REPORT_PATH.exists():
        shutil.copyfile(REPORT_PATH, BACKUP_PATH)
    REPORT_PATH.write_text(markdown, encoding="utf-8")


def main() -> None:
    jobs = load_jobs()

    reactant_job = jobs["reactant"].job
    scan_job = jobs["scan"].job
    ts_job = jobs["ts"].job
    product_job = jobs["product"].job

    scan_results = scan_job.results.get_pesscan_results()
    scan_energies = np.array(scan_results["PES"], dtype=float)
    scan_distances_bohr = np.array(scan_results["RaveledPESCoords"][0], dtype=float)
    scan_distances = scan_distances_bohr * Units.conversion_factor("bohr", "angstrom")
    scan_molecules = scan_results["Molecules"]

    highest_index = int(np.argmax(scan_energies))
    shortest_index = int(np.argmin(scan_distances))
    longest_index = int(np.argmax(scan_distances))

    energy_plot_path = Path("scan_energy_vs_distance.png")
    shortest_image_path = Path("scan_shortest_distance.png")
    longest_image_path = Path("scan_longest_distance.png")
    highest_image_path = Path("scan_highest_energy.png")
    reactant_image_path = Path("reactant_optimized.png")
    ts_image_path = Path("ts_optimized.png")
    product_image_path = Path("product_optimized.png")

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(scan_distances, scan_energies, marker="o")
    ax.set_xlabel("O...C distance (Angstrom)")
    ax.set_ylabel("Energy (hartree)")
    ax.set_title("GFN1-xTB PES scan")
    fig.tight_layout()
    fig.savefig(energy_plot_path, dpi=200)
    plt.close(fig)

    render_system_image(scan_molecules[shortest_index], shortest_image_path)
    render_system_image(scan_molecules[longest_index], longest_image_path)
    render_system_image(scan_molecules[highest_index], highest_image_path)
    render_system_image(reactant_job.results.get_main_system(), reactant_image_path)
    render_system_image(ts_job.results.get_main_system(), ts_image_path)
    render_system_image(product_job.results.get_main_system(), product_image_path)

    electronic = {
        "Reactant": energy_hartree(reactant_job),
        "TS": energy_hartree(ts_job),
        "Product": energy_hartree(product_job),
    }
    gibbs = {
        "Reactant": gibbs_hartree(reactant_job),
        "TS": gibbs_hartree(ts_job),
        "Product": gibbs_hartree(product_job),
    }
    rel_electronic = relative_kcal(electronic, "Reactant")
    rel_gibbs = relative_kcal(gibbs, "Reactant")

    summary = pd.DataFrame(
        {
            "Electronic Energy (Eh)": electronic,
            "Relative Electronic (kcal/mol)": rel_electronic,
            "Gibbs Free Energy (Eh)": gibbs,
            "Relative Gibbs (kcal/mol)": rel_gibbs,
        }
    )
    summary.index.name = "Stationary Point"

    ts_frequencies = np.array(ts_job.results.get_frequencies(unit="cm^-1"), dtype=float)
    imaginary_frequencies = ts_frequencies[ts_frequencies < 0.0]

    barrier_g_kcal = rel_gibbs["TS"]
    reaction_free_energy_kcal = rel_gibbs["Product"]

    markdown = f"""# Claisen Rearrangement with GFN1-xTB

This report summarizes a GFN1-xTB workflow for the Claisen rearrangement of `{SMILES}`. The workflow consisted of reactant optimization with normal modes, a one-dimensional PES scan of the carbonyl oxygen to terminal vinyl carbon distance, a transition-state search started from the highest-energy scan point, and a product optimization from the shortest-distance scan structure.

## Conclusions

- Reaction free energy barrier from Gibbs free energies: `{barrier_g_kcal:.2f} kcal/mol`
- Reaction free energy from Gibbs free energies: `{reaction_free_energy_kcal:.2f} kcal/mol`
- Transition-state verification: `{len(imaginary_frequencies)}` imaginary frequency/frequencies
- Imaginary frequency values (cm^-1): `{", ".join(f"{value:.2f}" for value in imaginary_frequencies)}`

## Scan Profile

![Energy vs distance]({energy_plot_path})

The PES scan used the directly identified carbonyl oxygen and terminal vinyl carbon and sampled the O...C distance across `{len(scan_distances)}` points, ending at `{scan_distances[shortest_index]:.3f} Angstrom`.

### Representative scan structures

Shortest-distance structure (`{scan_distances[shortest_index]:.3f} Angstrom`, `{scan_energies[shortest_index]:.6f} Eh`)

![Shortest scan structure]({shortest_image_path})

Longest-distance structure (`{scan_distances[longest_index]:.3f} Angstrom`, `{scan_energies[longest_index]:.6f} Eh`)

![Longest scan structure]({longest_image_path})

Highest-energy structure (`{scan_distances[highest_index]:.3f} Angstrom`, `{scan_energies[highest_index]:.6f} Eh`)

![Highest-energy scan structure]({highest_image_path})

## Optimized stationary points

{summary.to_markdown(floatfmt=".6f")}

Optimized reactant

![Reactant]({reactant_image_path})

Optimized transition state

![Transition state]({ts_image_path})

Optimized product

![Product]({product_image_path})

## Job provenance

- `reactant_opt`: {jobs["reactant"].why}
- `pes_scan`: {jobs["scan"].why}
- `ts_search`: {jobs["ts"].why}
- `product_opt`: {jobs["product"].why}

## Calculation inputs

### `reactant_opt`

```text
{reactant_job.get_input()}
```

### `pes_scan`

```text
{scan_job.get_input()}
```

### `ts_search`

```text
{ts_job.get_input()}
```

### `product_opt`

```text
{product_job.get_input()}
```
"""

    write_report(markdown)


if __name__ == "__main__":
    main()