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

"""Plot the mobile-region RDF before and after a GST124 melt--quench run."""

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scm.plams import AMSJob, AMSRDFJob, Settings, finish, init


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

    md = AMSJob.load_external(str(args.results.resolve()))
    molecule = md.results.get_input_molecule()
    mobile_atoms = [i for i, atom in enumerate(molecule, 1) if "mobile" in atom.properties.region]
    if not mobile_atoms:
        parser.error("The input structure has no atoms labeled region=mobile")

    nframes = int(md.results.readrkf("History", "nEntries"))
    frame_spacing_fs = md.results.get_time_step()
    window_frames = round(20_000 / frame_spacing_fs)  # final 20 ps
    if nframes <= window_frames + 1:
        parser.error("The trajectory is too short to compare its first frame with its last 20 ps")

    init(folder="mobile_rdf_workdir")
    try:
        curves = []
        for name, first, last in (
            ("Initial ordered geometry", 1, 1),
            ("Final 20 ps at 300 K", nframes - window_frames, nframes),
        ):
            settings = Settings()
            settings.input.TrajectoryInfo.Trajectory = [Settings()]
            settings.input.TrajectoryInfo.Trajectory[0].Range = f"{first} {last}"
            rdf = AMSRDFJob(
                previous_job=md,
                atom_indices=mobile_atoms,
                atom_indices_to=mobile_atoms,
                rmin=2.0,
                rmax=8.5,
                rstep=0.05,
                settings=settings,
                name=f"rdf_{first}_{last}",
            )
            rdf.run()
            if not rdf.check():
                raise RuntimeError(f"RDF analysis failed for {name}")
            curves.append((name, *rdf.results.get_rdf()))
    finally:
        finish()

    fig, axes = plt.subplots(2, 1, figsize=(8, 5.5), sharex=True)
    for ax, (name, r, g), color in zip(axes, curves, ("#2878B5", "#D55E00")):
        ax.plot(r, g, color=color, linewidth=1.8)
        ax.set_title(name, loc="left", fontsize=11)
        ax.set_ylabel("Mobile–mobile $g(r)$")
        ax.set_xlim(2.0, 8.5)
        ax.set_ylim(bottom=0)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
    axes[-1].set_xlabel("Distance (Å)")
    fig.tight_layout()
    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=220)
    plt.close(fig)
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()