#!/usr/bin/env python3
"""
Radial Density Profile (RDP) for micelles — companion to cmc.py
CMM 5/2026

For each frame, identifies aggregates >= min_size molecules, unwraps each
aggregate across PBC, then builds a radial number-density profile of every
bead type from the aggregate center of mass.  Also computes per-frame mean
radius of gyration and the three principal moments of inertia.

Results averaged over all qualifying aggregates across all frames.

Usage: python rdp.py id dcdFreq [ion1 [ion2]] [options]
  ion1, ion2 are optional; if omitted all non-W residues are used.
  --workers N          parallel worker processes (default: all CPUs)
  --min-size N         minimum aggregate size in molecules (default: 10)
  --dr FLOAT           radial bin width in Angstrom (default: 1.0)
  --rmax FLOAT         maximum radius in Angstrom (default: 60.0)
  --itp FILE [FILE ...] ITP/TOP file(s) to read Martini bead types from;
                        groups atoms by type (e.g. C2A+C3A+C4A → C1)
                        rather than by atom name

Output:
  RDP-{suff}/radial_density.dat  — number density (A^-3) vs r, per bead type
  RDP-{suff}/shape.dat           — time series: t  Rg  I1  I2  I3 per frame
"""

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

# ---------------------------------------------------------------------------
# ITP/TOP parser — extracts (resname, atomname) -> bead_type from [atoms]
# ---------------------------------------------------------------------------

def parse_bead_types(itp_paths):
    """Read one or more Martini ITP/TOP files and return a dict mapping
    (resname, atomname) -> bead_type from every [atoms] section found.
    Lines starting with ';' and blank lines are skipped.
    #include directives are NOT followed — list all relevant ITP files explicitly.
    """
    mapping = {}
    for path in itp_paths:
        in_atoms = False
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith(';'):
                    continue
                if line.startswith('['):
                    in_atoms = 'atoms' in line.lower()
                    continue
                if in_atoms:
                    parts = line.split()
                    # [atoms] columns: nr  type  resnr  residue  atom  cgnr  charge  mass
                    if len(parts) >= 5 and parts[0].isdigit():
                        bead_type = parts[1]
                        resname   = parts[3]
                        atomname  = parts[4]
                        mapping[(resname, atomname)] = bead_type
    return mapping


# ---------------------------------------------------------------------------
# Union-Find clustering (same as cmc.py)
# ---------------------------------------------------------------------------

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:
            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 — module-level for pickling by spawn workers
# ---------------------------------------------------------------------------

def _rdp_frame(args):
    (positions, box,
     atom_to_res, res_indices,
     res_to_atoms, atom_masses, atom_label_idx, n_labels,
     thresh_dist, min_size, bin_edges) = args

    from MDAnalysis.lib.distances import capped_distance
    from collections import defaultdict
    import numpy as np

    n_bins = len(bin_edges) - 1
    counts     = np.zeros((n_bins, n_labels), dtype=np.float64)
    rg_sum     = 0.0
    moments_sum = np.zeros(3, dtype=np.float64)   # I1, I2, I3 summed over aggregates
    n_analyzed = 0

    # --- cluster finding (identical to cmc.py) ---
    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)

    for cluster in cluster_list:
        if len(cluster) < min_size:
            continue

        # Gather atom data for this cluster
        atom_idx     = np.concatenate([res_to_atoms[ri] for ri in cluster])
        clust_pos    = positions[atom_idx].copy()
        clust_masses = atom_masses[atom_idx]
        clust_labels = atom_label_idx[atom_idx]

        # Unwrap across PBC: shift each atom to its minimum-image position
        # relative to the first atom so the cluster is contiguous.
        ref   = clust_pos[0]
        delta = clust_pos - ref
        for dim in range(3):
            L = box[dim]
            delta[:, dim] -= L * np.round(delta[:, dim] / L)
        clust_pos = ref + delta

        # Mass-weighted center of mass
        total_mass = clust_masses.sum()
        com = (clust_pos * clust_masses[:, np.newaxis]).sum(axis=0) / total_mass

        # Positions relative to COM
        r_vec = clust_pos - com                          # (N, 3)
        r_sq  = np.sum(r_vec ** 2, axis=1)              # |r_i|^2 for each atom

        # Radius of gyration: Rg = sqrt(sum(m_i * |r_i|^2) / M)
        rg = np.sqrt(np.dot(clust_masses, r_sq) / total_mass)
        rg_sum += rg

        # Inertia tensor (3x3 symmetric):
        # I_ab = sum_i m_i (|r_i|^2 delta_ab - r_ia * r_ib)
        inertia = (np.dot(clust_masses, r_sq) * np.eye(3)
                   - np.einsum('k,ki,kj->ij', clust_masses, r_vec, r_vec))
        # Principal moments = eigenvalues, sorted ascending (I1 <= I2 <= I3)
        moments_sum += np.linalg.eigvalsh(inertia)

        # Radial distance of each atom from COM for density profile
        r = np.sqrt(r_sq)
        bin_idx = np.searchsorted(bin_edges[1:], r)
        for bi, li in zip(bin_idx, clust_labels):
            if bi < n_bins:
                counts[bi, li] += 1

        n_analyzed += 1

    return n_analyzed, counts, rg_sum, moments_sum


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Micelle radial density profile")
    parser.add_argument("sim_id")
    parser.add_argument("dcd_freq", type=int)
    parser.add_argument("ion1", nargs="?", default=None,
                        help="residue name to analyze (default: all non-W)")
    parser.add_argument("ion2", nargs="?", default=None)
    parser.add_argument("--workers",  type=int,   default=multiprocessing.cpu_count())
    parser.add_argument("--min-size", type=int,   default=10,
                        help="minimum aggregate size in molecules (default: 10)")
    parser.add_argument("--dr",       type=float, default=1.0,
                        help="radial bin width in Angstrom (default: 1.0)")
    parser.add_argument("--rmax",     type=float, default=60.0,
                        help="maximum radius in Angstrom (default: 60.0)")
    parser.add_argument("--itp",      nargs="+",  default=[],  metavar="FILE",
                        help="ITP/TOP file(s) to read bead types from")
    pargs = parser.parse_args()

    sim_id    = pargs.sim_id
    dcd_freq  = pargs.dcd_freq
    ion1      = pargs.ion1
    ion2      = pargs.ion2
    n_workers = pargs.workers
    min_size  = pargs.min_size
    dr        = pargs.dr
    rmax      = pargs.rmax
    itp_files = pargs.itp

    # Built-in grouping for TTA/BzO wet and dry Martini atom names.
    # Maps (resname, atomname) -> display label; ITP files override these.
    BUILTIN_GROUPS = {
        # TTA
        ("TTA", "NC3"): "NC3",
        ("TTA", "C1A"): "SC1",   # neck bead
        ("TTA", "C2A"): "C1",    # tail beads grouped as C1
        ("TTA", "C3A"): "C1",
        ("TTA", "C4A"): "C1",
        # BzO / BZO — ring beads grouped as Rn
        ("BzO", "R1"):  "Rn",
        ("BzO", "R2"):  "Rn",
        ("BzO", "R3"):  "Rn",
        ("BZO", "R1"):  "Rn",
        ("BZO", "R2"):  "Rn",
        ("BZO", "R3"):  "Rn",
        ("BzO", "O4"):  "Qcar",
        ("BZO", "O4"):  "Qcar",
        ("BzO", "COO"): "Qcar",
        ("BZO", "COO"): "Qcar",
        # Dry Martini SC4 ring names (if used)
        ("BzO", "SC4"): "Rn",
        ("BZO", "SC4"): "Rn",
    }

    bead_type_map = parse_bead_types(itp_files) if itp_files else {}
    if bead_type_map:
        print(f"Loaded bead types from: {itp_files}")
    else:
        bead_type_map = BUILTIN_GROUPS
        print("Using built-in bead-type grouping (pass --itp to override)")

    thresh_dist = 10.0
    stride      = 1

    bin_edges = np.arange(0.0, rmax + dr, dr)
    bin_mids  = 0.5 * (bin_edges[:-1] + bin_edges[1:])
    n_bins    = len(bin_mids)

    # -----------------------------------------------------------------------
    # Load topology (before selection so we can auto-detect resnames)
    # -----------------------------------------------------------------------
    print("Loading 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")

    # Build selection — auto-detect if no residue names given
    if ion1 is None:
        resnames = sorted({r.resname for r in u.residues if r.resname != "W"})
        select_text = " or ".join(f"resname {r}" for r in resnames)
        suff = "-".join(resnames)
        print(f"Auto-detected residues: {resnames}")
    elif ion2 is None:
        select_text = f"resname {ion1}"
        suff = ion1
    else:
        select_text = f"resname {ion1} or resname {ion2}"
        suff = f"{ion1}-{ion2}"

    out_dir = f"RDP-{suff}"

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

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

    # Bead-type labels: "RESNAME-beadtype" when ITP provided, else "RESNAME-atomname"
    if bead_type_map:
        raw_labels = [
            f"{a.resname}-{bead_type_map.get((a.resname, a.name), a.name)}"
            for a in sel.atoms
        ]
    else:
        raw_labels = [f"{a.resname}-{a.name}" for a in sel.atoms]
    unique_labels = sorted(set(raw_labels))
    label_to_idx  = {lab: i for i, lab in enumerate(unique_labels)}
    atom_label_idx = np.array([label_to_idx[lab] for lab in raw_labels])
    n_labels       = len(unique_labels)
    print(f"Bead types: {unique_labels}")

    # Residue -> atom indices within the selection
    res_to_atoms_d = defaultdict(list)
    for atom_idx, ri in enumerate(atom_to_res.tolist()):
        res_to_atoms_d[ri].append(atom_idx)
    res_to_atoms = {ri: np.array(idxs) for ri, idxs in res_to_atoms_d.items()}

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

    total_counts    = np.zeros((n_bins, n_labels), dtype=np.float64)
    total_analyzed  = 0
    shape_rows      = []   # (t, mean_Rg, mean_I1, mean_I2, mean_I3) per frame

    WINDOW  = n_workers * 4
    pending = deque()   # (t, future) in submission order

    with ProcessPoolExecutor(max_workers=n_workers) as executor:
        for _fi, ts in enumerate(u.trajectory[::stride]):
            t = ts.time * 1e-3   # ps -> ns
            future = executor.submit(
                _rdp_frame,
                (sel.positions.copy(), ts.dimensions.copy(),
                 atom_to_res, res_indices,
                 res_to_atoms, atom_masses, atom_label_idx, n_labels,
                 thresh_dist, min_size, bin_edges)
            )
            pending.append((t, future))

            if len(pending) >= WINDOW:
                t_r, fut_r = pending.popleft()
                n_a, counts, rg_sum, moments_sum = fut_r.result()
                total_counts   += counts
                total_analyzed += n_a
                if n_a > 0:
                    shape_rows.append((t_r, rg_sum / n_a,
                                       moments_sum[0] / n_a,
                                       moments_sum[1] / n_a,
                                       moments_sum[2] / n_a))

        while pending:
            t_r, fut_r = pending.popleft()
            n_a, counts, rg_sum, moments_sum = fut_r.result()
            total_counts   += counts
            total_analyzed += n_a
            if n_a > 0:
                shape_rows.append((t_r, rg_sum / n_a,
                                   moments_sum[0] / n_a,
                                   moments_sum[1] / n_a,
                                   moments_sum[2] / n_a))

    print(f"Analyzed {total_analyzed} aggregates (Cn >= {min_size})")

    # -----------------------------------------------------------------------
    # Write radial_density.dat
    # -----------------------------------------------------------------------
    shell_vol = 4.0 * np.pi * bin_mids**2 * dr
    shell_vol[shell_vol == 0] = np.nan

    density = (total_counts / (shell_vol[:, np.newaxis] * total_analyzed)
               if total_analyzed > 0 else np.zeros_like(total_counts))

    rdp_path = f"{out_dir}/radial_density.dat"
    with open(rdp_path, "w") as f:
        f.write(f"# Radial density profile — {sim_id} ({select_text})\n")
        f.write(f"# Aggregates analyzed (Cn >= {min_size}): {total_analyzed}\n")
        f.write(f"# Normalization: number density per aggregate (A^-3)\n")
        f.write(f"# Columns: r_mid(A)  " + "  ".join(unique_labels) + "\n")
        for ri in range(n_bins):
            row = "  ".join(f"{density[ri, li]:12.6e}" for li in range(n_labels))
            f.write(f"{bin_mids[ri]:8.3f}  {row}\n")
    print(f"Written: {rdp_path}")

    # -----------------------------------------------------------------------
    # Write shape.dat  (per-frame time series)
    # -----------------------------------------------------------------------
    shape_path = f"{out_dir}/shape.dat"
    with open(shape_path, "w") as f:
        f.write(f"# Shape parameters — {sim_id} ({select_text})\n")
        f.write(f"# Per-frame means over aggregates with Cn >= {min_size}\n")
        f.write(f"# Rg in Angstrom; I1<=I2<=I3 in amu*A^2\n")
        f.write(f"# t(ns)      Rg(A)       I1          I2          I3\n")
        for row in shape_rows:
            f.write(f"{row[0]:10.4f}  {row[1]:10.4f}  "
                    f"{row[2]:12.4f}  {row[3]:12.4f}  {row[4]:12.4f}\n")

    if shape_rows:
        arr = np.array(shape_rows)
        print(f"Written: {shape_path}")
        print(f"  Mean Rg = {arr[:,1].mean():.3f} A  "
              f"  Mean I1/I2/I3 = {arr[:,2].mean():.1f} / "
              f"{arr[:,3].mean():.1f} / {arr[:,4].mean():.1f} amu*A^2")
