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

Usage: python cmc.py id ion1 [ion2]
  ion2 is optional; omitting it selects only ion1.
  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 numpy as np
import MDAnalysis as mda
from MDAnalysis.lib.distances import capped_distance

# ---------------------------------------------------------------------------
# Arguments (mirrors TCL: id, ion1, ion2 optional)
# ---------------------------------------------------------------------------
if len(sys.argv) < 3:
    print("Usage: python cmc.py id ion1 [ion2]")
    sys.exit(1)

sim_id = sys.argv[1]
ion1   = sys.argv[2]
ion2   = sys.argv[3] if len(sys.argv) > 3 else ""

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   # Angstroms; for Martini, 8-12 A gives same result
dcd_freq    = 1000   # nstxout-compressed
timestep    = 30     # fs (dt)
stride      = 1

frame_interval = 1e-6 * timestep * dcd_freq  # ns per frame
if stride < 1:
    stride = 1
dt = frame_interval * stride  # ns between analysed frames

# ---------------------------------------------------------------------------
# Cluster utilities (direct port of TCL procs)
# ---------------------------------------------------------------------------

def check_overlap(a_list, b_list):
    b_set = set(b_list)
    return any(a in b_set for a in a_list)


def search_overlap(nll):
    """One pass: find and merge the first overlapping pair."""
    for i in range(len(nll) - 1):
        for j in range(i + 1, len(nll)):
            if check_overlap(nll[i], nll[j]):
                nll[i] = sorted(set(nll[i]) | set(nll[j]))
                nll.pop(j)
                return nll
    return nll


def get_clusters(nll):
    """Iteratively merge overlapping neighbour lists into unique clusters."""
    n, n0 = len(nll), 0
    while n != n0:
        n0 = n
        nll = search_overlap(nll)
        n = len(nll)
    return nll


# ---------------------------------------------------------------------------
# Load topology and precompute stable per-residue data
# ---------------------------------------------------------------------------
print("Loading topology...")
# Use directory name for the TPR (mirrors TCL: "set cwd [file tail [pwd]]")
# so one TPR serves all id variants (1.0mM-run-vis.tpr, 1.0mM-last100ns-run-vis.xtc, etc.)
cwd = os.path.basename(os.getcwd())
u   = mda.Universe(f"{cwd}-run-vis.tpr", 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}'")

# Map sel-atom position index -> resindex (stable across frames)
atom_to_res = np.array([a.resindex for a in sel.atoms])

# Per-residue total mass (all atoms in residue, matching VMD "atomselect top residue N")
resindex_to_mass = {}
for ri in res_indices:
    resindex_to_mass[ri] = float(u.atoms[u.atoms.resindices == ri].masses.sum())

# ---------------------------------------------------------------------------
# Stage 1 — Trajectory loop
# ---------------------------------------------------------------------------
print("Reading the trajectory...")
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()
total_time     = 0.0
n_frames_total = 0

for fi, ts in enumerate(u.trajectory[::stride]):
    t = fi * dt

    positions = sel.positions   # (N_atoms, 3) — updates each step
    box       = ts.dimensions   # periodic box for PBC-aware distances

    # Single capped_distance call replaces per-residue "pbwithin" selections
    pairs, _ = capped_distance(positions, positions, thresh_dist, box=box)

    # Build per-residue neighbour sets; each residue starts as its own neighbour
    neigh = {ri: {ri} for ri in res_indices}
    for ai, aj in pairs:
        ri, rj = int(atom_to_res[ai]), int(atom_to_res[aj])
        neigh[ri].add(rj)
        neigh[rj].add(ri)

    neigh_list_list = [sorted(s) for s in neigh.values()]
    cluster_list    = get_clusters(neigh_list_list)
    cluster_num     = len(cluster_list)

    if cluster_num == 0:
        print("ERROR: Found zero clusters. Problem with the selection text?")
        sys.exit(1)

    size_list       = []
    mass_list_frame = []
    for cluster in cluster_list:
        size_list.append(len(cluster))
        cluster_mass = sum(resindex_to_mass[ri] for ri in cluster)
        mass_list_frame.append(cluster_mass)
        mass_list_master.add(cluster_mass)

    size_list.sort()
    mass_list_frame.sort()

    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")

    total_time      = t
    n_frames_total += 1

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 = {}  # mc[lineno][iclust] = count
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 = {}  # iclust -> cumulative count across all frames

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):
            ts_val = counter * timestep * 0.001
            count  = mc[counter].get(iclust, 0)
            out_f.write(f"{ts_val} {count}\n")
            sum_mass[iclust] += count

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

# ---------------------------------------------------------------------------
# Stage 3 — Summations (old stage3.tcl)
# ---------------------------------------------------------------------------
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]))
x1           = sum_mass[monomer_mass] * normalization

with open(f"{data_dir}/Sizes.dat", "w") as out2, \
     open(f"{data_dir}/Masses.dat", "w") as out3:

    for fh, label in [(out2, "Sizes -- frac, no. vs Cn for"),
                      (out3, "Masses -- frac, no. vs mass for")]:
        fh.write(f"# {label}\n")
        fh.write(f"# {sim_id} {ion1} {ion2}\n")
        fh.write(f"# {total_time} ns total time\n")
        fh.write(f"# X1: {x1}\n")
        fh.write(f"# sum Xn: {1.0 - x1}\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")
        out2.write(f"{iclust // monomer_mass:<5d} {val * normalization:<7.4f} {val:<7.4f}\n")

print("Finished with summations.")
