#!/usr/bin/env amspython

from __future__ import annotations

from pathlib import Path

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


TRAJECTORY = Path("GUI_tests/ReaxFF_Production.results")
REGION = "resin"
BINS = 80
START_FRACTION = 0.5
OUTPUT_PREFIX = "DENSMAP_resin"


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_indices(molecule, region: str) -> list[int]:
    return [i for i, atom in enumerate(molecule, start=1) if atom_region(atom) == region]


def density_map_xy(
    job: AMSJob,
    region: str,
    bins: int,
    start_fraction: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    n_frames = job.results.readrkf("History", "nEntries")
    start_frame = max(1, min(n_frames, int(n_frames * start_fraction) + 1))

    final_molecule = job.results.get_main_molecule()
    lx = final_molecule.lattice[0][0]
    ly = final_molecule.lattice[1][1]
    histogram = np.zeros((bins, bins), dtype=float)

    log(f"Using frames {start_frame}-{n_frames} of {n_frames}")
    log(f"Density region: {region}")

    for used_frames, frame in enumerate(range(start_frame, n_frames + 1), start=1):
        molecule = job.results.get_history_molecule(frame)
        atoms = list(molecule)
        indices = region_indices(molecule, region)
        if not indices:
            raise ValueError(f"No atoms found for region '{region}'")

        coords = np.asarray([atoms[i - 1].coords[:2] for i in indices], dtype=float)
        coords[:, 0] %= lx
        coords[:, 1] %= ly
        histogram += np.histogram2d(
            coords[:, 0],
            coords[:, 1],
            bins=bins,
            range=[[0.0, lx], [0.0, ly]],
        )[0]

        if used_frames == 1:
            log(f"Atoms in region: {len(indices)}")
            log(f"Map size: {lx / 10.0:.3f} x {ly / 10.0:.3f} nm")
        if frame == n_frames or used_frames % 10 == 0:
            log(f"Processed {used_frames}/{n_frames - start_frame + 1} frames")

    dx_nm = (lx / bins) / 10.0
    dy_nm = (ly / bins) / 10.0
    density = histogram / used_frames / (dx_nm * dy_nm)
    x_edges = np.linspace(0.0, lx / 10.0, bins + 1)
    y_edges = np.linspace(0.0, ly / 10.0, bins + 1)
    return x_edges, y_edges, density.T


def main() -> None:
    log(f"Loading trajectory: {TRAJECTORY}")
    job = AMSJob.load_external(TRAJECTORY)
    x_edges, y_edges, density = density_map_xy(job, REGION, BINS, START_FRACTION)

    np.savetxt(f"{OUTPUT_PREFIX}.txt", density, header="resin atoms / nm^2 / frame")

    fig, ax = plt.subplots(figsize=(5.0, 4.0))
    mesh = ax.pcolormesh(x_edges, y_edges, density, shading="auto")
    fig.colorbar(mesh, ax=ax, label="resin atoms / nm2 / frame")
    ax.set_xlabel("x (nm)")
    ax.set_ylabel("y (nm)")
    ax.set_title("Resin density map")
    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()