#!/usr/bin/env python3
"""
CMC Analysis - Python port of cmc.tcl
Derived from TDMD Analysis, CMM 5/2026

Usage: python cmc.py id dcdFreq ion1 [ion2] [--workers N]
  dcdFreq:    nstxout-compressed value (e.g. 1000 or 5000 for long runs).
  ion2:       optional; omitting it selects only ion1.
  --workers N: parallel worker processes (default: all CPUs).
  The TPR is taken from the current directory name so one TPR
  can serve id variants (e.g. 1.0mM-last100ns uses 1.0mM-run-vis.tpr).
"""

import sys
import os
import argparse
import multiprocessing
from collections import defaultdict, deque
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import MDAnalysis as mda
from MDAnalysis.lib.distances import capped_distance

# ---------------------------------------------------------------------------
# Clustering — Union-Find  O(N α(N)) vs O(N²) for the original TCL algorithm
# ---------------------------------------------------------------------------

def union_find_clusters(res_indices, neigh_pairs):
    parent = {ri: ri for ri in res_indices}
    rank   = {ri: 0  for ri in res_indices}

    def find(x):
        root = x
        while parent[root] != root:
            root = parent[root]
        while parent[x] != root:          # path compression
            parent[x], x = root, parent[x]
        return root

    def union(x, y):
        rx, ry = find(x), find(y)
        if rx == ry:
            return
        if rank[rx] < rank[ry]:
            rx, ry = ry, rx
        parent[ry] = rx
        if rank[rx] == rank[ry]:
            rank[rx] += 1

    for ri, rj in neigh_pairs:
        union(ri, rj)

    clusters = defaultdict(list)
    for ri in res_indices:
        clusters[find(ri)].append(ri)
    return list(clusters.values())


# ---------------------------------------------------------------------------
# Per-frame worker — must be module-level so it can be pickled by spawn workers
# ---------------------------------------------------------------------------

def _analyze_frame(args):
    positions, box, atom_to_res, res_indices, resindex_to_mass, thresh_dist = args
    from MDAnalysis.lib.distances import capped_distance
    from collections import defaultdict

    pairs, _ = capped_distance(positions, positions, thresh_dist, box=box)

    neigh_pairs = set()
    for ai, aj in pairs:
        ri, rj = int(atom_to_res[ai]), int(atom_to_res[aj])
        if ri != rj:
            neigh_pairs.add((ri, rj))

    cluster_list = union_find_clusters(res_indices, neigh_pairs)
    cluster_num  = len(cluster_list)

    size_list       = []
    mass_list_frame = []
    cluster_masses  = []
    for cluster in cluster_list:
        size_list.append(len(cluster))
        cm = sum(resindex_to_mass[ri] for ri in cluster)
        mass_list_frame.append(cm)
        cluster_masses.append(cm)

    size_list.sort()
    mass_list_frame.sort()
    return cluster_num, size_list, mass_list_frame, cluster_masses


# ---------------------------------------------------------------------------
# Main — guarded so spawn workers don't re-execute it on import
# ---------------------------------------------------------------------------

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="CMC cluster analysis")
    parser.add_argument("sim_id")
    parser.add_argument("dcd_freq", type=int)
    parser.add_argument("ion1")
    parser.add_argument("ion2", nargs="?", default="")
    parser.add_argument("--workers", type=int, default=multiprocessing.cpu_count(),
                        help="parallel worker processes (default: all CPUs)")
    pargs = parser.parse_args()

    sim_id    = pargs.sim_id
    dcd_freq  = pargs.dcd_freq
    ion1      = pargs.ion1
    ion2      = pargs.ion2
    n_workers = pargs.workers

    if not ion2:
        select_text = f"resname {ion1}"
        suff = ion1
    else:
        select_text = f"resname {ion1} or resname {ion2}"
        suff = f"{ion1}-{ion2}"

    out_dir  = f"Cluster-{suff}"
    data_dir = f"Data-{suff}"

    # -----------------------------------------------------------------------
    # MD parameters
    # -----------------------------------------------------------------------
    thresh_dist = 10.0
    # dcd_freq comes from argv; long runs use 5000+ (5x or 10x default)
    # timestep is used only in stage 2 mass-file timestamps; stage 1 reads
    # actual time directly from ts.time (ps) in the trajectory metadata.
    timestep = 30    # fs (dt) — fallback only; override with --timestep if needed
    stride   = 1

    # -----------------------------------------------------------------------
    # Load topology
    # -----------------------------------------------------------------------
    print("Loading topology...")
    # Use directory name for the topology (mirrors TCL: "set cwd [file tail [pwd]]")
    # so one topology file serves all id variants.
    # Try TPR first; fall back to GRO if the TPR version is too new for MDAnalysis
    # (MDAnalysis <2.8 does not support tpx v138 / GROMACS 2025+).
    # Note: unlike VMD, MDAnalysis residue selection works correctly with GRO topology.
    cwd = os.path.basename(os.getcwd())
    try:
        u = mda.Universe(f"{cwd}-run-vis.tpr", f"{sim_id}-run-vis.xtc")
    except (ValueError, NotImplementedError) as e:
        print(f"TPR load failed ({e}); falling back to GRO topology.")
        u = mda.Universe(f"{cwd}-run-vis.gro", f"{sim_id}-run-vis.xtc")

    sel          = u.select_atoms(select_text)
    res_indices  = sorted({a.resindex for a in sel.atoms})
    num_residues = len(res_indices)
    print(f"Found {num_residues} residues matching '{select_text}'")

    atom_to_res = np.array([a.resindex for a in sel.atoms])

    resindex_to_mass = {
        ri: float(u.atoms[u.atoms.resindices == ri].masses.sum())
        for ri in res_indices
    }

    # -----------------------------------------------------------------------
    # Stage 1 — Parallel trajectory loop
    # -----------------------------------------------------------------------
    print(f"Reading the trajectory (workers={n_workers})...")
    os.makedirs(out_dir, exist_ok=True)

    f_mean = open(f"{out_dir}/clusterSizeMean.dat", "w")
    f_num  = open(f"{out_dir}/clusterNum.dat",      "w")
    f_mass = open(f"{out_dir}/clusterMass.dat",     "w")
    f_dist = open(f"{out_dir}/clusterDist.dat",     "w")
    f_max  = open(f"{out_dir}/clusterSizeMax.dat",  "w")

    mass_list_master = set()
    size_hist        = defaultdict(int)   # cluster size (molecule count) -> total occurrences over all frames
    frame_times      = []   # actual time (ns) per frame, used in stage 2
    total_time       = 0.0
    n_frames_total   = 0

    # Sliding window: keep n_workers*4 futures in flight without loading the
    # whole trajectory into memory.  A single for-loop over the trajectory
    # avoids the while-True+break iterator pattern that caused 4x repetition.
    WINDOW = n_workers * 4
    pending = deque()   # (t, future) in submission order

    def write_result(t, future):
        cluster_num, size_list, mass_list_frame, cluster_masses = future.result()
        if cluster_num == 0:
            print("ERROR: Found zero clusters. Problem with the selection text?")
            sys.exit(1)
        mass_list_master.update(cluster_masses)
        for s in size_list:              # size histogram from actual molecule counts
            size_hist[s] += 1
        f_dist.write(" ".join(str(s) for s in size_list) + "\n")
        f_mean.write(f"{t} {sum(size_list) / len(size_list)}\n")
        f_max.write( f"{t} {size_list[-1]}\n")
        f_num.write( f"{t:<8.4f} {cluster_num}\n")
        f_mass.write(f"{t:9.4f} {' '.join(str(m) for m in mass_list_frame)}\n")
        frame_times.append(t)

    with ProcessPoolExecutor(max_workers=n_workers) as executor:
        for _fi, ts in enumerate(u.trajectory[::stride]):
            # ts.time is in ps; convert to ns.  This uses the actual trajectory
            # timestamps rather than a computed dt, so timestep/dcdFreq errors
            # can't corrupt the time axis.
            t = ts.time * 1e-3
            future = executor.submit(
                _analyze_frame,
                (sel.positions.copy(), ts.dimensions.copy(),
                 atom_to_res, res_indices, resindex_to_mass, thresh_dist)
            )
            pending.append((t, future))
            # Drain oldest result once the window is full
            if len(pending) >= WINDOW:
                write_result(*pending.popleft())
        # Drain remainder
        while pending:
            write_result(*pending.popleft())

    total_time     = frame_times[-1] if frame_times else 0.0
    n_frames_total = len(frame_times)

    for fh in (f_mean, f_num, f_mass, f_max, f_dist):
        fh.close()

    mass_list_master = sorted(mass_list_master)

    # -----------------------------------------------------------------------
    # Stage 2 — Histogram compilation
    # -----------------------------------------------------------------------
    print("Compiling histograms...")

    mc = {}
    with open(f"{out_dir}/clusterMass.dat") as infile_m:
        for lineno, line in enumerate(infile_m):
            mc[lineno] = {round(c): 0 for c in mass_list_master}
            for token in line.split():
                iclust = round(float(token))
                mc[lineno][iclust] = mc[lineno].get(iclust, 0) + 1

    l_max    = lineno
    sum_mass = {}

    for clust in mass_list_master:
        iclust = round(clust)
        sum_mass[iclust] = 0
        with open(f"{out_dir}/mass{iclust}.dat", "w") as out_f:
            for counter in range(1, l_max + 1):
                count  = mc[counter].get(iclust, 0)
                out_f.write(f"{frame_times[counter]} {count}\n")
                sum_mass[iclust] += count

    print(f"Masses found: {mass_list_master}")

    # -----------------------------------------------------------------------
    # Stage 3 — Summations
    # -----------------------------------------------------------------------
    print("Writing output...")

    normalization = 1.0 / (n_frames_total * num_residues)
    os.makedirs(data_dir, exist_ok=True)

    monomer_mass = int(round(mass_list_master[0]))

    # ---- Size histogram (Sizes.dat): built from actual molecule counts, ----
    # ---- independent of masses so it is robust to a mixed-species        ----
    # ---- selection (TTA+BzO) and to guessed masses on GRO fallback.      ----
    # val[n] = n * (occurrences of size-n clusters) = molecules found in
    # size-n clusters summed over frames; frac = val * normalization = Xn.
    size_x1     = size_hist.get(1, 0) * 1 * normalization
    val_mono_sz = size_hist.get(1, 0)                       # monomer-cluster count (nsize=1)
    agg_ge10_sz = sum(n * cnt for n, cnt in size_hist.items() if n >= 10)
    x_param_sz  = agg_ge10_sz / val_mono_sz if val_mono_sz > 0 else 0.0

    # MSF of aggregation number: sqrt(<n²> - <n>²) / N
    # averaged over all individual cluster sizes across all frames
    total_clusters = sum(size_hist.values())
    mean_n  = sum(n * cnt for n, cnt in size_hist.items()) / total_clusters
    mean_n2 = sum(n * n * cnt for n, cnt in size_hist.items()) / total_clusters
    msf_n   = (mean_n2 - mean_n ** 2) ** 0.5 / num_residues

    with open(f"{data_dir}/Sizes.dat", "w") as out2:
        out2.write("# Sizes -- frac, no. vs Cn for\n")
        out2.write(f"# {sim_id} {ion1} {ion2}\n")
        out2.write(f"# {total_time} ns total time\n")
        out2.write(f"# X1: {size_x1}\n")
        out2.write(f"# sum Xn: {1.0 - size_x1}\n")
        out2.write(f"# x (agg>=10 / monomer): {x_param_sz}\n")
        out2.write(f"# MSF(n)/N: {msf_n}\n")
        for n in sorted(size_hist):
            val = n * size_hist[n]
            out2.write(f"{n:<5d} {val * normalization:<7.4f} {val:<7.4f}\n")

    # ---- Mass histogram (Masses.dat) kept as-is for reference ----
    x1 = sum_mass[monomer_mass] * normalization
    val_monomer = sum_mass[monomer_mass]
    agg_ge10    = sum(
        (int(round(clust)) // monomer_mass) * sum_mass[int(round(clust))]
        for clust in mass_list_master
        if int(round(clust)) // monomer_mass >= 10
    )
    x_param = agg_ge10 / val_monomer if val_monomer > 0 else 0.0

    with open(f"{data_dir}/Masses.dat", "w") as out3:
        out3.write("# Masses -- frac, no. vs mass for\n")
        out3.write(f"# {sim_id} {ion1} {ion2}\n")
        out3.write(f"# {total_time} ns total time\n")
        out3.write(f"# X1: {x1}\n")
        out3.write(f"# sum Xn: {1.0 - x1}\n")
        out3.write(f"# x (agg>=10 / monomer): {x_param}\n")

        for clust in mass_list_master:
            iclust = int(round(clust))
            nsize  = iclust // monomer_mass
            val    = nsize * sum_mass[iclust]

            with open(f"{data_dir}/sumMass_{iclust}.dat", "w") as out1:
                out1.write(f"{clust:<7.2f} {sum_mass[iclust]:<5d}\n")

            out3.write(f"{iclust:<5d} {val * normalization:<7.4f} {val:<7.4f}\n")

    print("Finished with summations.")
