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

"""Plot near-90-degree angles from one AMS MD results folder.

Run: amspython 01-plot-near-90-angles.py path/to/job.results
The script reads ams.rkf and writes one PNG figure; no other files are needed.
"""

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scm.plams import AMSJob


CUTOFF_ANGSTROM = 3.6
ANGLE_RANGE_DEG = (80.0, 100.0)
SAMPLE_EVERY_PS = 10.0
FIT_RANGE_PS = (50.0, 200.0)


def near_90_count(molecule, mobile: np.ndarray) -> int:
    """Count 80--100 degree angles with a mobile center and two mobile neighbors."""
    xyz = np.asarray([atom.coords for atom in molecule.atoms], dtype=float)[mobile]
    cell = np.asarray(molecule.lattice, dtype=float)
    if cell.shape != (3, 3):
        raise ValueError("A three-dimensional periodic cell is required")
    inverse_cell = np.linalg.inv(cell)
    count = 0
    for center in xyz:
        vectors = xyz - center
        fractional = vectors @ inverse_cell
        vectors = (fractional - np.rint(fractional)) @ cell
        distances = np.linalg.norm(vectors, axis=1)
        neighbors = vectors[(distances > 1e-8) & (distances < CUTOFF_ANGSTROM)]
        if len(neighbors) < 2:
            continue
        unit = neighbors / np.linalg.norm(neighbors, axis=1)[:, None]
        a, b = np.triu_indices(len(unit), k=1)
        angles = np.degrees(np.arccos(np.clip(np.sum(unit[a] * unit[b], axis=1), -1, 1)))
        count += int(np.count_nonzero((angles >= ANGLE_RANGE_DEG[0]) & (angles <= ANGLE_RANGE_DEG[1])))
    return count


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("results", type=Path, help="AMS MD job.results folder containing ams.rkf")
    parser.add_argument("--output", type=Path, default=Path("near_90_angles.png"))
    args = parser.parse_args()
    if not (args.results / "ams.rkf").is_file():
        parser.error(f"No ams.rkf in {args.results}")

    job = AMSJob.load_external(str(args.results.resolve()))
    molecule = job.results.get_input_molecule()
    mobile = np.asarray([i for i, atom in enumerate(molecule.atoms)
                         if "mobile" in atom.properties.get("region", [])], dtype=int)
    if len(mobile) == 0:
        parser.error("The input structure has no atoms labeled region=mobile")

    times = np.asarray(job.results.get_history_property("Time", history_section="MDHistory"), dtype=float) / 1000
    times -= times[0]
    targets = np.arange(0, times[-1] + SAMPLE_EVERY_PS / 2, SAMPLE_EVERY_PS)
    frames = sorted({int(np.argmin(abs(times - target))) + 1 for target in targets})
    sampled_times = times[np.asarray(frames) - 1]
    counts = np.asarray([near_90_count(job.results.get_history_molecule(frame), mobile)
                         for frame in frames], dtype=float)
    normalized = counts * 1000 / len(mobile)

    fit = (sampled_times >= FIT_RANGE_PS[0]) & (sampled_times <= FIT_RANGE_PS[1])
    if np.count_nonzero(fit) < 3:
        raise ValueError(f"Trajectory does not cover the {FIT_RANGE_PS[0]:g}--{FIT_RANGE_PS[1]:g} ps fit window")
    slope, intercept = np.polyfit(sampled_times[fit], normalized[fit], 1)
    rate = slope * 1000  # ps to ns

    output = args.output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    fig, ax = plt.subplots(figsize=(7.0, 4.3))
    ax.plot(sampled_times, normalized, "o", color="#2878B5", markersize=4)
    ax.plot(sampled_times[fit], slope * sampled_times[fit] + intercept, "--",
            color="#172B4D", linewidth=2,
            label=f"50–200 ps fit: {rate:,.0f} angles / (1,000 mobile atoms·ns)")
    ax.set(xlabel="Time (ps)", ylabel="80–100° angles per 1,000 mobile atoms",
           title="Near-90° angular order during recrystallization", xlim=(0, times[-1]))
    ax.legend(frameon=False, loc="upper left")
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    fig.savefig(output, dpi=220)
    plt.close(fig)
    print(f"Saved {output}; fitted rate: {rate:.1f} angles/(1,000 mobile atoms·ns)")


if __name__ == "__main__":
    main()