#!/usr/bin/env amspython

from __future__ import annotations

from pathlib import Path
import warnings

import matplotlib.pyplot as plt
import numpy as np
from scm.plams import AMSJob, Molecule


TRAJECTORY = Path("GUI_tests/ReaxFF_Production.results")
MOLECULE_FILES = {
    "polymer": Path("PI_30mer.xyz"),
    "resin": Path("HDCPD_trimer.xyz"),
}
RDF_PAIRS = (("resin", "resin"), ("resin", "polymer"), ("polymer", "polymer"))
RMAX = 40.0
DR = 0.1
START_FRACTION = 0.5
OUTPUT_PREFIX = "RDF_com_regions"


def log(message: str) -> None:
    print(message, flush=True)


def atom_region(atom) -> str | None:
    properties = atom.properties
    return properties.get("region") if hasattr(properties, "get") else getattr(properties, "region", None)


def region_molecules(molecule: Molecule, region: str, atoms_per_molecule: int) -> list[list[int]]:
    atom_indices = [i for i, atom in enumerate(molecule, start=1) if atom_region(atom) == region]
    if not atom_indices:
        raise ValueError(f"No atoms found for region '{region}'")
    if len(atom_indices) % atoms_per_molecule:
        raise ValueError(
            f"Region '{region}' has {len(atom_indices)} atoms, which is not divisible by "
            f"{atoms_per_molecule} atoms per molecule"
        )
    return [
        atom_indices[i : i + atoms_per_molecule]
        for i in range(0, len(atom_indices), atoms_per_molecule)
    ]


def lattice_matrix(molecule: Molecule) -> np.ndarray:
    return np.asarray(molecule.lattice, dtype=float)


def minimum_image_vectors(vectors: np.ndarray, lattice: np.ndarray) -> np.ndarray:
    fractional = vectors @ np.linalg.inv(lattice)
    return vectors - np.round(fractional) @ lattice


def center_of_mass(molecule: Molecule, atom_indices: list[int]) -> np.ndarray:
    molecule_atoms = list(molecule)
    atoms = [molecule_atoms[i - 1] for i in atom_indices]
    masses = np.asarray([atom.mass for atom in atoms], dtype=float)
    coords = np.asarray([atom.coords for atom in atoms], dtype=float)
    lattice = lattice_matrix(molecule)
    unwrapped = coords[0] + minimum_image_vectors(coords - coords[0], lattice)
    com = np.average(unwrapped, axis=0, weights=masses)
    return com - np.floor(com @ np.linalg.inv(lattice)) @ lattice


def pair_distances(points_a: np.ndarray, points_b: np.ndarray, lattice: np.ndarray, same_region: bool) -> np.ndarray:
    if same_region:
        vectors = [
            points_a[j] - points_a[i]
            for i in range(len(points_a))
            for j in range(i + 1, len(points_a))
        ]
    else:
        vectors = [point_b - point_a for point_a in points_a for point_b in points_b]
    if not vectors:
        return np.empty(0, dtype=float)
    return np.linalg.norm(minimum_image_vectors(np.asarray(vectors), lattice), axis=1)


def calculate_rdfs(
    job: AMSJob,
    pairs: tuple[tuple[str, str], ...],
    atoms_per_molecule: dict[str, int],
    rmax: float,
    dr: float,
    start_fraction: float,
) -> tuple[np.ndarray, dict[str, np.ndarray]]:
    n_frames = job.results.readrkf("History", "nEntries")
    start_frame = max(1, min(n_frames, int(n_frames * start_fraction) + 1))
    first_molecule = job.results.get_history_molecule(start_frame)
    max_r = min(np.linalg.norm(vector) for vector in lattice_matrix(first_molecule)) / 2.0
    if rmax > max_r:
        warnings.warn(f"Capping rmax from {rmax:.3f} to {max_r:.3f} A", RuntimeWarning)
        rmax = max_r

    bins = np.arange(0.0, rmax + dr, dr)
    histograms = {f"{a}-{b}": np.zeros(len(bins) - 1, dtype=float) for a, b in pairs}
    volumes: list[float] = []
    molecule_counts: dict[str, int] | None = None

    log(f"Using frames {start_frame}-{n_frames} of {n_frames}")
    log("RDF pairs: " + ", ".join(f"{a}-{b}" for a, b in pairs))

    for used_frames, frame in enumerate(range(start_frame, n_frames + 1), start=1):
        molecule = job.results.get_history_molecule(frame)
        centers = {
            region: np.asarray(
                [
                    center_of_mass(molecule, group)
                    for group in region_molecules(molecule, region, atoms_per_molecule[region])
                ]
            )
            for region in atoms_per_molecule
        }

        counts = {region: len(points) for region, points in centers.items()}
        if molecule_counts is None:
            molecule_counts = counts
            log("Molecule counts: " + ", ".join(f"{k}={v}" for k, v in counts.items()))
        elif counts != molecule_counts:
            raise ValueError(f"Frame {frame} has a different number of molecules")

        lattice = lattice_matrix(molecule)
        for region_a, region_b in pairs:
            distances = pair_distances(
                centers[region_a], centers[region_b], lattice, region_a == region_b
            )
            histograms[f"{region_a}-{region_b}"] += np.histogram(distances, bins=bins)[0]
        volumes.append(molecule.unit_cell_volume())

        if frame == n_frames or used_frames % 10 == 0:
            log(f"Processed {used_frames}/{n_frames - start_frame + 1} frames")

    if molecule_counts is None or not volumes:
        raise ValueError("No trajectory frames were used")

    radii = 0.5 * (bins[:-1] + bins[1:])
    shell_volumes = 4.0 * np.pi / 3.0 * (bins[1:] ** 3 - bins[:-1] ** 3)
    volume = float(np.mean(volumes))

    rdf_data: dict[str, np.ndarray] = {}
    for region_a, region_b in pairs:
        n_a = molecule_counts[region_a]
        n_b = molecule_counts[region_b]
        pair_count = n_a * (n_a - 1) / 2.0 if region_a == region_b else n_a * n_b
        normalization = len(volumes) * shell_volumes * pair_count / volume
        rdf_data[f"{region_a}-{region_b}"] = histograms[f"{region_a}-{region_b}"] / normalization

    return radii, rdf_data


def main() -> None:
    log(f"Loading trajectory: {TRAJECTORY}")
    job = AMSJob.load_external(TRAJECTORY)
    atoms_per_molecule = {
        region: len(Molecule(filename))
        for region, filename in MOLECULE_FILES.items()
    }

    radii, rdf_data = calculate_rdfs(
        job, RDF_PAIRS, atoms_per_molecule, RMAX, DR, START_FRACTION
    )

    fig, ax = plt.subplots(figsize=(5.2, 3.4))
    for label, values in rdf_data.items():
        ax.plot(radii, values, label=label)

    table = np.column_stack([radii, *rdf_data.values()])
    header = "r_A " + " ".join(rdf_data)
    np.savetxt(f"{OUTPUT_PREFIX}.txt", table, header=header)

    ax.set_xlabel("r (A)")
    ax.set_ylabel("g(r)")
    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(f"{OUTPUT_PREFIX}.png")
    log(f"Wrote {OUTPUT_PREFIX}.txt")
    log(f"Wrote {OUTPUT_PREFIX}.png")


if __name__ == "__main__":
    main()