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

from dataclasses import dataclass
from pathlib import Path

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


SMILES = "c1ccccc1-c2ccccc2"
JOB_NAME = "biphenyl_dftb_torsion"


@dataclass(frozen=True)
class TorsionDefinition:
    central_bond: tuple[int, int]
    dihedral: tuple[int, int, int, int]

    @property
    def central_bond_one_based(self) -> tuple[int, int]:
        return tuple(idx + 1 for idx in self.central_bond)

    @property
    def dihedral_one_based(self) -> tuple[int, int, int, int]:
        return tuple(idx + 1 for idx in self.dihedral)


def carbon_neighbors(system: ChemicalSystem, atom_index: int, exclude: int) -> list[int]:
    return [
        idx
        for idx in system.bonds.get_bonded_atoms(atom_index)
        if idx != exclude and system.atoms[idx].symbol == "C"
    ]


def identify_inter_ring_torsion(system: ChemicalSystem) -> TorsionDefinition:
    candidates: list[tuple[int, int]] = []
    for i, j, _bond in system.bonds:
        if system.atoms[i].symbol == "C" and system.atoms[j].symbol == "C":
            if system.bond_cuts_molecule(i, j):
                candidates.append((i, j))

    if len(candidates) != 1:
        raise RuntimeError(f"Expected one central inter-ring C-C bond, found {candidates}")

    left, right = candidates[0]
    left_neighbors = carbon_neighbors(system, left, exclude=right)
    right_neighbors = carbon_neighbors(system, right, exclude=left)
    if not left_neighbors or not right_neighbors:
        raise RuntimeError("Could not find ring-neighbor carbon atoms for the inter-ring dihedral")

    return TorsionDefinition(
        central_bond=(left, right),
        dihedral=(left_neighbors[0], left, right, right_neighbors[0]),
    )


def latest_job_path() -> Path:
    candidates = sorted(Path(".").glob("01-run-biphenyl-dftb-pesscan_workdir*/" + JOB_NAME))
    if not candidates:
        raise FileNotFoundError("No biphenyl PES scan job directory was found")
    return candidates[-1]


def pesscan_dataframe(job: AMSJob) -> tuple[pd.DataFrame, list[ChemicalSystem]]:
    results = job.results.get_pesscan_results()
    angles = pd.Series(results["RaveledPESCoords"][0], dtype=float)
    if angles.abs().max() <= 2.0 * 3.141592653589793:
        angles = angles * Units.conversion_factor("radian", "degree")

    energies_hartree = pd.Series(results["PES"], dtype=float)
    min_energy_hartree = energies_hartree.min()
    relative_kcal = (energies_hartree - min_energy_hartree) * Units.conversion_factor("Hartree", "kcal/mol")

    df = pd.DataFrame(
        {
            "Dihedral angle (deg)": angles.round(6),
            "Absolute energy (Hartree)": energies_hartree,
            "Relative energy (kcal/mol)": relative_kcal,
        }
    )
    return df, list(results["Molecules"])


def nearest_index(df: pd.DataFrame, target: float) -> int:
    return int((df["Dihedral angle (deg)"] - target).abs().idxmin())


def write_plot(df: pd.DataFrame) -> None:
    fig, ax = plt.subplots(figsize=(6.5, 4.2))
    ax.plot(df["Dihedral angle (deg)"], df["Relative energy (kcal/mol)"], marker="o", linewidth=1.8)
    ax.set_xlabel("Dihedral angle (deg)")
    ax.set_ylabel("Relative energy (kcal/mol)")
    ax.set_xlim(0, 180)
    ax.grid(True, alpha=0.25)
    fig.tight_layout()
    fig.savefig("biphenyl_torsion_profile.png", dpi=180)
    plt.close(fig)


def write_structure_images(df: pd.DataFrame, molecules: list[ChemicalSystem]) -> dict[str, str]:
    min_idx = int(df["Relative energy (kcal/mol)"].idxmin())
    image_specs = {
        "planar": nearest_index(df, 0.0),
        "minimum": min_idx,
        "perpendicular": nearest_index(df, 90.0),
    }
    output: dict[str, str] = {}
    for label, idx in image_specs.items():
        filename = f"biphenyl_{label}.png"
        view(molecules[idx], width=360, height=280, guess_bonds=True, direction="along_pca3", picture_path=filename)
        output[label] = filename
    return output


def write_report(job: AMSJob, df: pd.DataFrame, images: dict[str, str], torsion: TorsionDefinition) -> None:
    min_idx = int(df["Relative energy (kcal/mol)"].idxmin())
    planar0_idx = nearest_index(df, 0.0)
    planar180_idx = nearest_index(df, 180.0)
    perpendicular_idx = nearest_index(df, 90.0)

    min_angle = df.loc[min_idx, "Dihedral angle (deg)"]
    min_twist = min(min_angle, 180.0 - min_angle)
    planar_barrier = min(
        df.loc[planar0_idx, "Relative energy (kcal/mol)"],
        df.loc[planar180_idx, "Relative energy (kcal/mol)"],
    )
    perpendicular_barrier = df.loc[perpendicular_idx, "Relative energy (kcal/mol)"]

    table = df.copy()
    table["Absolute energy (Hartree)"] = table["Absolute energy (Hartree)"].map(lambda x: f"{x:.10f}")
    table["Relative energy (kcal/mol)"] = table["Relative energy (kcal/mol)"].map(lambda x: f"{x:.4f}")
    table["Dihedral angle (deg)"] = table["Dihedral angle (deg)"].map(lambda x: f"{x:.1f}")

    report = f"""# Biphenyl DFTB torsional PES scan

## Setup

Biphenyl was built from SMILES `{SMILES}` with `ChemicalSystem.from_smiles`. The central inter-ring C-C bond was identified by inspecting `ChemicalSystem.bonds` for the carbon-carbon bond whose removal cuts the molecule.

- Central inter-ring bond, zero-based atom indices: `{torsion.central_bond}`
- Central inter-ring bond, AMS one-based atom indices: `{torsion.central_bond_one_based}`
- Inter-ring dihedral, zero-based atom indices: `{torsion.dihedral}`
- Inter-ring dihedral, AMS one-based atom indices: `{torsion.dihedral_one_based}`

The constrained AMS `PESScan` used DFTB with `Model GFN1-xTB`, scanning the inter-ring dihedral from 0 to 180 degrees in 19 points, i.e. 10 degree spacing. Energies below are reported relative to the lowest scan point.

## Results

![Biphenyl torsional profile](biphenyl_torsion_profile.png)

Minimum-energy scan point: **{min_angle:.1f} degrees**.

Minimum-energy biphenyl twist angle: **{min_twist:.1f} degrees**. This is reported as `min(dihedral, 180 - dihedral)` because the 0 to 180 degree scan has equivalent twisted structures on either side of 90 degrees.

Torsional barrier to the planar structure: **{planar_barrier:.4f} kcal/mol**. The 0 and 180 degree planar endpoints are both tabulated below; the quoted planar barrier is the lower of the two endpoint relative energies.

Torsional barrier to the perpendicular structure: **{perpendicular_barrier:.4f} kcal/mol**.

## Representative optimized structures

Planar geometry:

![Planar biphenyl]({images["planar"]})

Minimum-energy geometry:

![Minimum-energy biphenyl]({images["minimum"]})

Perpendicular geometry:

![Perpendicular biphenyl]({images["perpendicular"]})

## Scan table

{table.to_markdown(index=False, disable_numparse=True)}

## AMS input

The PES scan job was loaded from `{job.path}` with `AMSJob.load_external(...)`.

```ams
{job.get_input()}
```

## Conclusion

The DFTB scan places the minimum at a scan dihedral of {min_angle:.1f} degrees, corresponding to a biphenyl twist of {min_twist:.1f} degrees. Relative to this point, the lower planar endpoint is {planar_barrier:.4f} kcal/mol higher and the perpendicular point is {perpendicular_barrier:.4f} kcal/mol higher.
"""
    Path("report.md").write_text(report)


def main() -> None:
    system = ChemicalSystem.from_smiles(SMILES)
    torsion = identify_inter_ring_torsion(system)

    job = AMSJob.load_external(str(latest_job_path()))
    df, molecules = pesscan_dataframe(job)
    write_plot(df)
    images = write_structure_images(df, molecules)
    write_report(job, df, images, torsion)


if __name__ == "__main__":
    main()