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

import re
from pathlib import Path
from typing import Any, Dict, List

import jinja2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MaxNLocator
from scm.base import Units
from scm.plams import AMSJob, view


WORKDIR_STEM = "01-run_workdir"
JOB_NAME = "ethane_uff_optimization"
TEMPLATE_PATH = Path("report.template.md.j2")


def natural_key(path: Path) -> List[Any]:
    return [
        int(part) if part.isdigit() else part
        for part in re.split(r"(\d+)", str(path))
    ]


def latest_workdir(stem: str) -> Path:
    candidates = [path for path in Path(".").glob("{}*".format(stem)) if path.is_dir()]
    if not candidates:
        raise FileNotFoundError("No PLAMS work directory found for {}".format(stem))
    return max(candidates, key=natural_key)


def carbon_indices(molecule: Any) -> List[int]:
    indices = [
        index
        for index, atom in enumerate(molecule.atoms)
        if atom.symbol == "C"
    ]
    if len(indices) != 2:
        raise ValueError("Expected exactly two carbon atoms, found {}".format(len(indices)))
    return indices


def cc_distance(molecule: Any) -> float:
    first, second = carbon_indices(molecule)
    first_coords = np.asarray(molecule.atoms[first].coords, dtype=float)
    second_coords = np.asarray(molecule.atoms[second].coords, dtype=float)
    return float(np.linalg.norm(first_coords - second_coords))


def render_report() -> None:
    workdir = latest_workdir(WORKDIR_STEM)
    job_path = workdir / JOB_NAME
    job = AMSJob.load_external(str(job_path))
    job_description = (
        "This job is the sole calculation used in the report. It optimizes a "
        "coordinate-perturbed ethane structure with the AMS ForceField engine "
        "and the UFF parameter set."
    )

    history_length = job.results.get_history_length()
    energies_hartree = np.asarray(
        job.results.get_history_property("Energy"), dtype=float
    )
    if len(energies_hartree) != history_length:
        raise ValueError("Energy history length does not match geometry history")

    hartree_to_kcal_mol = Units.conversion_factor("hartree", "kcal/mol")
    energies_kcal_mol = energies_hartree * hartree_to_kcal_mol
    molecules = [
        job.results.get_history_molecule(step)
        for step in range(1, history_length + 1)
    ]
    if any(molecule is None for molecule in molecules):
        raise ValueError("A geometry could not be read from the AMS history")
    cc_lengths = np.asarray([cc_distance(molecule) for molecule in molecules])
    steps = np.arange(1, history_length + 1)

    tables_dir = Path("tables")
    figures_dir = Path("figures")
    tables_dir.mkdir(exist_ok=True)
    figures_dir.mkdir(exist_ok=True)

    trajectory = pd.DataFrame(
        {
            "Step": steps,
            "Energy (kcal/mol)": energies_kcal_mol,
            "C-C distance (angstrom)": cc_lengths,
        }
    )
    trajectory.to_csv(tables_dir / "optimization_history.csv", index=False)

    summary = pd.DataFrame(
        [
            {
                "Optimization steps": history_length,
                "Initial energy (kcal/mol)": energies_kcal_mol[0],
                "Final energy (kcal/mol)": energies_kcal_mol[-1],
                "Energy change (kcal/mol)": energies_kcal_mol[-1]
                - energies_kcal_mol[0],
                "Initial C-C distance (angstrom)": cc_lengths[0],
                "Final C-C distance (angstrom)": cc_lengths[-1],
            }
        ]
    )
    summary.to_csv(tables_dir / "optimization_summary.csv", index=False)

    fig, ax = plt.subplots(figsize=(6.4, 4.0))
    ax.plot(steps, energies_kcal_mol, marker="o", linewidth=1.5)
    ax.set_xlabel("Optimization step")
    ax.set_ylabel("Energy (kcal/mol)")
    ax.xaxis.set_major_locator(MaxNLocator(nbins=9, integer=True))
    ax.grid(alpha=0.25)
    fig.tight_layout()
    energy_path = figures_dir / "energy_vs_step.png"
    fig.savefig(str(energy_path), dpi=180)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(6.4, 4.0))
    ax.plot(steps, cc_lengths, marker="o", linewidth=1.5)
    ax.set_xlabel("Optimization step")
    ax.set_ylabel("C-C distance (angstrom)")
    ax.xaxis.set_major_locator(MaxNLocator(nbins=9, integer=True))
    ax.grid(alpha=0.25)
    fig.tight_layout()
    distance_path = figures_dir / "cc_distance_vs_step.png"
    fig.savefig(str(distance_path), dpi=180)
    plt.close(fig)

    optimized_system = job.results.get_main_system()
    structure_path = figures_dir / "optimized_ethane.png"
    view(
        optimized_system,
        guess_bonds=len(optimized_system.bonds) == 0,
        direction="along_pca3",
        width=600,
        height=450,
        picture_path=str(structure_path),
    )

    timings = job.results.get_timings()
    elapsed_time = float(timings.get("elapsed", timings.get("total", 0.0)))
    energy_change = energies_kcal_mol[-1] - energies_kcal_mol[0]
    conclusion = (
        "The UFF optimization completed in {} recorded geometry steps. The "
        "energy changed by {:.3f} kcal/mol, and the C-C distance changed from "
        "{:.4f} to {:.4f} angstrom."
    ).format(history_length, energy_change, cc_lengths[0], cc_lengths[-1])
    method = (
        "Ethane was generated from the SMILES string `CC` with "
        "`ChemicalSystem.from_smiles`. `ChemicalSystem.perturb_coordinates` "
        "then added a random displacement in the closed interval from -0.1 to "
        "+0.1 angstrom to every Cartesian coordinate. AMS optimized this "
        "structure with `Task GeometryOptimization` and the ForceField engine "
        "using UFF. Energies and structures were read from the AMS History "
        "section. Energy conversion used `scm.base.Units.conversion_factor`."
    )

    context: Dict[str, Any] = {
        "title": "UFF geometry optimization of ethane",
        "purpose": (
            "Optimize a perturbed ethane geometry with UFF and track the "
            "energy and carbon-carbon distance through the optimization."
        ),
        "conclusion": conclusion,
        "figures": [
            {
                "alt": "Energy versus geometry optimization step",
                "path": str(energy_path),
                "caption": "UFF energy at each recorded geometry step.",
            },
            {
                "alt": "Carbon-carbon distance versus geometry optimization step",
                "path": str(distance_path),
                "caption": "Carbon-carbon distance at each recorded geometry step.",
            },
            {
                "alt": "Optimized ethane structure",
                "path": str(structure_path),
                "caption": "Final ethane geometry from the UFF optimization.",
            },
        ],
        "tables": [
            {
                "caption": (
                    "Optimization summary. The complete trajectory is available "
                    "in `tables/optimization_history.csv`."
                ),
                "df_markdown": summary.to_markdown(
                    index=False,
                    floatfmt=(".0f", ".6f", ".6f", ".6f", ".6f", ".6f"),
                ),
            }
        ],
        "method": method,
        "jobs": [
            {
                "name": job.name,
                "path": str(job_path),
                "elapsed_time": elapsed_time,
                "description": job_description,
                "input": job.get_input().strip(),
            }
        ],
    }

    template = jinja2.Environment(
        loader=jinja2.FileSystemLoader(str(TEMPLATE_PATH.parent)),
        autoescape=False,
        keep_trailing_newline=True,
    ).get_template(TEMPLATE_PATH.name)
    Path("report.md").write_text(template.render(**context), encoding="utf-8")


if __name__ == "__main__":
    render_report()