#!/usr/bin/env python3
"""
Single-aggregate form factor P(q) by direct Debye summation.

    P(q) = (1/N^2) SUM_i SUM_j  sin(q r_ij) / (q r_ij)

computed over the beads of ONE micelle, so the only distances involved are
INTRA-aggregate. This sidesteps the box/truncation limit that cripples the
g(r)->S(q) route (sq.py): there, r_max = L_box/2 sets a lowest trustworthy
q ~ 2*pi/r_max, which for these systems (0.5-0.8 nm^-1) sits at or above
1/R of the aggregate -- exactly where sphere and cylinder form factors
start to differ. Hence sphere and cylinder fits to S(q) look identical.
Here the accessible q range is set by the micelle's own extent instead.

What this gives / does not give:
  + a clean single-aggregate form factor: the q^-1 rod regime, the Holtzer
    plateau, and the cross-sectional Guinier are all visible and usable,
    so sphere-vs-cylinder is actually decidable.
  - NO inter-micelle structure factor (no correlation peak). This is P(q),
    not I(q) = P(q)S(q). For the correlation peak use sq.py.

Method per frame:
  1. cluster the surfactant by 10 A contact (same convention as cmc.py/beta.py)
  2. take the largest cluster (>= --min-size molecules)
  3. UNWRAP it across PBC by breadth-first search along the contact graph,
     shifting each molecule by the minimum image relative to its parent.
     (Molecules are already whole: make_vis.sh runs trjconv -pbc mol.)
  4. PERCOLATION CHECK: if the unwrapped cluster is inconsistent -- i.e. some
     contact edge disagrees with the minimum-image distance -- the aggregate
     wraps onto itself through the periodic boundary and its form factor is
     NOT well defined. Such frames are skipped and counted. (Expect this for
     the 200 mM box-spanning worm.)
  5. accumulate a fine-binned intra-aggregate distance histogram; P(q) is
     summed from the histogram at the end (far cheaper than an O(N^2) sum
     per q per frame, and accurate for bins << 2*pi/q_max).

Output: pq.dat with columns  q(nm^-1)  P(q)  q*P(q)   [Holtzer in col 3]

usage:
  pq.py SIM_ID SURF [--anion ABZO] [--stride N] [--max-frames N]
        [--qmin 0.05] [--qmax 10] [--nq 300] [--min-size 20]

then:
  fit_iq.py pq.dat --icol 1 --model cylinder --log
  fit_iq.py pq.dat --icol 1 --model sphere   --log
"""
import os
import sys
import argparse
from collections import defaultdict, deque

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


# ---------------------------------------------------------------------------
# clustering (same union-find / cutoff convention as cmc.py)
# ---------------------------------------------------------------------------
def union_find(res_indices, neigh_pairs):
    parent = {ri: ri 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

    for ri, rj in neigh_pairs:
        rx, ry = find(ri), find(rj)
        if rx != ry:
            parent[ry] = rx
    clusters = defaultdict(list)
    for ri in res_indices:
        clusters[find(ri)].append(ri)
    return list(clusters.values())


def unwrap_cluster(com, adj, members, box):
    """Unwrap molecule COMs of one cluster across PBC by BFS along contacts.

    com     : dict resid -> COM position (wrapped), nm
    adj     : dict resid -> set(neighbour resids) (contact graph)
    members : list of resids in this cluster
    box     : (3,) orthorhombic box lengths, nm

    Returns (unwrapped_com dict, shift dict, percolating bool).
    Percolation is detected when a contact edge between two already-placed
    molecules disagrees with the minimum-image convention: that means the
    cluster closes on itself through the boundary, so no consistent unwrap
    exists and P(q) is undefined.
    """
    L = np.asarray(box[:3], dtype=float)
    seed = members[0]
    unw = {seed: com[seed].copy()}
    shift = {seed: np.zeros(3)}
    order = deque([seed])
    placed = {seed}

    while order:
        u = order.popleft()
        for v in adj[u]:
            if v in placed:
                continue
            d = com[v] - com[u]
            d -= L * np.round(d / L)          # minimum image
            unw[v] = unw[u] + d
            shift[v] = unw[v] - com[v]
            placed.add(v)
            order.append(v)

    # consistency check over every contact edge
    percolating = False
    tol = 1e-6
    for u in members:
        for v in adj[u]:
            if v <= u:
                continue
            d_unw = unw[v] - unw[u]
            d_min = d_unw - L * np.round(d_unw / L)
            if np.any(np.abs(d_unw - d_min) > tol):
                percolating = True
                break
        if percolating:
            break

    return unw, shift, percolating


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Single-aggregate form factor P(q) by Debye summation")
    p.add_argument("sim_id")
    p.add_argument("surf", help="surfactant resname (e.g. TTA)")
    p.add_argument("--anion", default=None,
                   help="also include counterions bound to the aggregate (e.g. ABZO)")
    p.add_argument("--cutoff", type=float, default=10.0, help="contact cutoff, Angstrom (default 10)")
    p.add_argument("--min-size", type=int, default=20, help="min molecules for a cluster to count")
    p.add_argument("--stride", type=int, default=100, help="frame stride (default 100)")
    p.add_argument("--max-frames", type=int, default=200, help="max frames to average (default 200)")
    p.add_argument("--begin", type=float, default=None, help="start time, ns")
    p.add_argument("--end", type=float, default=None, help="end time, ns")
    p.add_argument("--qmin", type=float, default=0.05)
    p.add_argument("--qmax", type=float, default=10.0)
    p.add_argument("--nq", type=int, default=300)
    p.add_argument("--dr", type=float, default=0.02, help="distance histogram bin, nm")
    p.add_argument("-o", "--out", default="pq.dat")
    p.add_argument("--pr-out", default="pr.dat",
                   help="pair distance distribution p(r) output (default pr.dat)")
    p.add_argument("--axis", action="store_true",
                   help="REAL-SPACE cross-sectional radius via per-bead LOCAL CYLINDER AXIS. "
                        "Topology-independent: needs no unwrapping, no finite aggregate and "
                        "no dilute limit, so it works identically for spheres, worms, rings "
                        "and percolating networks -- unlike both the Debye route (needs an "
                        "unwrappable aggregate) and --local (needs a dilute-limit form "
                        "factor, which does not exist at 100/200 mM).")
    p.add_argument("--axis-nb", type=float, default=5.0,
                   help="neighbourhood radius (nm) for the local-axis PCA (default 5.0). "
                        "SENSITIVE PARAMETER: must exceed 2R (~4.7 nm here) so the slab "
                        "spans the whole cross-section, but larger values average over the "
                        "worm's curvature and inflate R. Calibrated on 50 mM against the "
                        "Debye truth R = 2.351 nm: nb = 4/5/6/7 give 2.20/2.43/2.53/2.57, "
                        "so 5.0 is the sweet spot (+3.5%). Re-check if R differs much from "
                        "~2.4 nm.")
    p.add_argument("--axis-slab", type=float, default=1.0,
                   help="thickness (nm) of the perpendicular slab used to locate the axis "
                        "(default 1.0). The slab spans the full cross-section, so its "
                        "centroid lies ON the axis -- a plain neighbourhood centroid does "
                        "not, and biases surface beads inward.")
    p.add_argument("--local", action="store_true",
                   help="PERCOLATING-SAFE mode for box-spanning aggregates (100/200 mM). "
                        "Computes S(q) EXACTLY by direct reciprocal-space summation over "
                        "the box's allowed q vectors -- no PBC unwrapping and no distance "
                        "truncation, so it works when the aggregate wraps onto itself and "
                        "the Debye route is undefined. Gives the cross-sectional radius "
                        "(a LOCAL quantity, unaffected by global topology). Does NOT give "
                        "N_agg, Rg, kappa^2 or contour length -- those are genuinely lost "
                        "under percolation.")
    p.add_argument("--qmax-local", type=float, default=3.0,
                   help="max q for --local reciprocal sum (nm^-1, default 3.0)")
    p.add_argument("--ring-test", action="store_true",
                   help="detect CLOSED-LOOP (ring/torus) aggregates. Shape metrics "
                        "(kappa^2, p(r), Guinier) CANNOT distinguish a closed ring from a "
                        "hairpin-folded worm -- that difference is topological, not "
                        "geometric. This projects onto the principal (ring) plane and "
                        "measures angular coverage about the centroid + whether the "
                        "centre is hollow.")
    p.add_argument("--guinier-qlo", type=float, default=None,
                   help="lower q for the cross-sectional Guinier fit (default 1.5*2pi/D_max)")
    p.add_argument("--guinier-qhi", type=float, default=None,
                   help="upper q for the cross-sectional Guinier fit (default: iterate to q*Rc=1.3)")
    a = p.parse_args()

    cwd = os.path.basename(os.getcwd())
    try:
        u = mda.Universe(f"{cwd}-run-vis.tpr", f"{a.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"{a.sim_id}-run-vis.xtc")

    surf = u.select_atoms(f"resname {a.surf}")
    if len(surf) == 0:
        sys.exit(f"no atoms with resname {a.surf}")
    anion = u.select_atoms(f"resname {a.anion}") if a.anion else None

    surf_res = sorted({at.resindex for at in surf.atoms})
    surf_a2r = np.array([at.resindex for at in surf.atoms])
    # map resindex -> row indices into surf.positions
    rows_of = defaultdict(list)
    for k, ri in enumerate(surf_a2r):
        rows_of[int(ri)].append(k)
    rows_of = {ri: np.asarray(v) for ri, v in rows_of.items()}

    if anion is not None:
        anion_a2r = np.array([at.resindex for at in anion.atoms])
        arows_of = defaultdict(list)
        for k, ri in enumerate(anion_a2r):
            arows_of[int(ri)].append(k)
        arows_of = {ri: np.asarray(v) for ri, v in arows_of.items()}

    # distance histogram accumulator
    nbins = None
    hist_acc = None
    rmax_seen = 0.0
    n_used = 0
    n_perc = 0
    n_small = 0
    nagg_list = []
    nbead_list = []
    rg_list = []
    shape_list = []   # (L1, L2, L3, asphericity, acylindricity, kappa^2)
    ring_list = []    # (t, angular coverage, core_fill, R_ring, is_ring, elongated)
    qvecs = qmags = qidx = qbins = sq_acc = sq_cnt = None   # --local reciprocal grid
    axis_list = []    # (R_from_<d^2>, <d>_headgroup, neighbourhood anisotropy)

    for ts in u.trajectory[::a.stride]:
        if n_used >= a.max_frames:
            break
        t = ts.time / 1000.0
        if a.begin is not None and t < a.begin:
            continue
        if a.end is not None and t > a.end:
            break

        box = ts.dimensions
        L_nm = np.asarray(box[:3], dtype=float) / 10.0

        # --- cluster surfactant by contact
        pairs = capped_distance(surf.positions, surf.positions, a.cutoff,
                                box=box, return_distances=False)
        adj = defaultdict(set)
        neigh = set()
        for ai, aj in pairs:
            ri, rj = int(surf_a2r[ai]), int(surf_a2r[aj])
            if ri != rj:
                neigh.add((ri, rj))
                adj[ri].add(rj)
                adj[rj].add(ri)
        clusters = union_find(surf_res, neigh)
        clusters.sort(key=len, reverse=True)
        big = clusters[0]
        if len(big) < a.min_size:
            n_small += 1
            continue

        # --- molecule COMs (nm)
        pos_nm = surf.positions / 10.0
        com = {ri: pos_nm[rows_of[ri]].mean(axis=0) for ri in big}

        # --- AXIS branch: real-space cross-section from per-bead local axes.
        # For each bead: take its neighbours within --axis-nb (minimum image, so
        # PBC and percolation are irrelevant), PCA their displacements, and call
        # the largest-eigenvalue direction the local cylinder axis. The bead's
        # perpendicular distance to the line through the neighbourhood centroid
        # along that axis is its cross-sectional radial coordinate d.
        #   uniform cylinder of radius R:  <d^2> = R^2/2  ->  R = sqrt(2<d^2>)
        # Headgroups (first bead of each surfactant) sit at the surface, so
        # <d>_head is an independent, assumption-free estimate of R.
        if a.axis:
            agg_rows = np.concatenate([rows_of[ri] for ri in big])
            is_head = np.concatenate([[True] + [False] * (len(rows_of[ri]) - 1)
                                      for ri in big])
            posA = surf.positions[agg_rows]
            boxA = np.asarray(box[:3], dtype=float)
            cut = a.axis_nb * 10.0
            prs = capped_distance(posA, posA, cut, box=box, return_distances=False)
            order = np.argsort(prs[:, 0], kind="stable")
            prs = prs[order]
            starts = np.searchsorted(prs[:, 0], np.arange(len(posA)))
            ends = np.searchsorted(prs[:, 0], np.arange(len(posA)) + 1)
            dperp = np.full(len(posA), np.nan)
            aniso = []
            for i in range(len(posA)):
                nb = prs[starts[i]:ends[i], 1]
                nb = nb[nb != i]
                if len(nb) < 8:
                    continue
                dd = posA[nb] - posA[i]
                dd -= boxA * np.round(dd / boxA)          # minimum image
                dd /= 10.0                                 # -> nm
                cen = dd.mean(axis=0)                      # centroid rel. to bead i
                dc = dd - cen
                C = (dc.T @ dc) / len(dc)
                w, V = np.linalg.eigh(C)
                u = V[:, 2]                                # local axis DIRECTION
                aniso.append(w[2] / max(w[0], 1e-12))
                # The full-ball centroid is NOT on the axis for a surface bead:
                # its neighbourhood is truncated by the aggregate surface, so the
                # centroid is pulled only part-way in and every distance comes out
                # too small. Instead locate the axis from a thin SLAB perpendicular
                # to u, which spans the whole local cross-section and whose centroid
                # therefore does sit on the axis.
                along = dd @ u
                slab = np.abs(along) < 0.5 * a.axis_slab
                if slab.sum() < 8:
                    continue
                c_ax = dd[slab].mean(axis=0)               # point on the axis
                v = -c_ax                                  # bead i relative to axis
                dperp[i] = np.linalg.norm(v - np.dot(v, u) * u)
            ok = np.isfinite(dperp)
            if ok.sum() > 50:
                axis_list.append((
                    float(np.sqrt(2.0 * np.mean(dperp[ok] ** 2))),      # R from <d^2>
                    float(np.mean(dperp[ok & is_head])) if (ok & is_head).any() else np.nan,
                    float(np.mean(aniso)),
                ))
            nagg_list.append(len(big))
            nbead_list.append(len(posA))
            n_used += 1
            continue

        # --- LOCAL (percolating-safe) branch: exact reciprocal-space S(q).
        # For a periodic box the scattering is exactly
        #     S(q) = <|SUM_j exp(i q.r_j)|^2> / N
        # evaluated on the allowed wavevectors q = 2*pi/L * (h,k,l), then
        # spherically averaged over each |q| shell. No unwrapping is required
        # (positions may stay wrapped) and nothing is truncated, so this stays
        # valid when the aggregate percolates and the Debye route does not exist.
        if a.local:
            agg_rows = np.concatenate([rows_of[ri] for ri in big])
            pos = pos_nm[agg_rows]
            Nl = len(pos)
            if qvecs is None:
                Hmax = int(np.ceil(a.qmax_local * float(np.max(L_nm)) / (2.0 * np.pi)))
                rng = np.arange(-Hmax, Hmax + 1)
                H, K, Lq = np.meshgrid(rng, rng, rng, indexing="ij")
                hkl = np.stack([H.ravel(), K.ravel(), Lq.ravel()], axis=1)
                hkl = hkl[np.any(hkl != 0, axis=1)]           # drop q=0
                qv = 2.0 * np.pi * hkl / L_nm                 # (Nq,3)
                qmag = np.linalg.norm(qv, axis=1)
                keep = qmag <= a.qmax_local
                qvecs, qmags = qv[keep], qmag[keep]
                # shell bins
                qbins = np.arange(0.0, a.qmax_local + 0.05, 0.05)
                qidx = np.digitize(qmags, qbins) - 1
                sq_acc = np.zeros(len(qbins) - 1)
                sq_cnt = np.zeros(len(qbins) - 1)
                for i in range(len(qbins) - 1):
                    sq_cnt[i] = np.sum(qidx == i)
            # |F(q)|^2 in chunks over q to bound memory
            Fsq = np.empty(len(qvecs))
            CH = 4000
            for s in range(0, len(qvecs), CH):
                ph = qvecs[s:s + CH] @ pos.T                  # (chunk, Nl)
                Fsq[s:s + CH] = (np.cos(ph).sum(axis=1) ** 2
                                 + np.sin(ph).sum(axis=1) ** 2)
            for i in range(len(sq_acc)):
                m = qidx == i
                if m.any():
                    sq_acc[i] += Fsq[m].mean() / Nl
            nagg_list.append(len(big))
            nbead_list.append(Nl)
            n_used += 1
            continue

        # --- unwrap across PBC, detect percolation
        unw, shift, percolating = unwrap_cluster(com, adj, big, L_nm)
        if percolating:
            n_perc += 1
            continue

        # --- gather unwrapped bead coordinates of the aggregate
        chunks = []
        for ri in big:
            chunks.append(pos_nm[rows_of[ri]] + shift[ri])
        if anion is not None:
            # bound anions: within cutoff of an aggregate surfactant bead.
            # place each bound anion by minimum image relative to the nearest
            # aggregate molecule COM, then apply that molecule's shift.
            apos_nm = anion.positions / 10.0
            agg_rows = np.concatenate([rows_of[ri] for ri in big])
            cp = capped_distance(anion.positions, surf.positions[agg_rows],
                                 a.cutoff, box=box, return_distances=False)
            bound = sorted({int(anion_a2r[ai]) for ai, _ in cp})
            if bound:
                # reference: nearest aggregate molecule (by COM) for each bound anion
                com_arr = np.array([com[ri] for ri in big])
                shf_arr = np.array([shift[ri] for ri in big])
                for ri_a in bound:
                    r_rows = arows_of[ri_a]
                    acom = apos_nm[r_rows].mean(axis=0)
                    d = com_arr - acom
                    d -= L_nm * np.round(d / L_nm)
                    k = int(np.argmin((d ** 2).sum(axis=1)))
                    # min-image the anion onto that molecule, then shift together
                    dd = acom - com_arr[k]
                    dd -= L_nm * np.round(dd / L_nm)
                    base = com_arr[k] + dd
                    chunks.append(apos_nm[r_rows] - acom + base + shf_arr[k])

        coords = np.concatenate(chunks, axis=0)
        N = len(coords)

        # --- gyration tensor: model-free shape, no form-factor assumption.
        # eigenvalues L1<=L2<=L3 of  S_ab = <r_a r_b>  about the centroid.
        #   Rg^2         = L1+L2+L3
        #   asphericity  b     = L3 - (L1+L2)/2        (0 for a sphere)
        #   acylindricity c    = L2 - L1               (0 for cylindrical symmetry)
        #   rel. shape anisotropy kappa^2 = (b^2 + 0.75 c^2)/Rg^4
        #     kappa^2 = 0 sphere, 0.25 disc, -> 1 for a thin rod.
        # This is the honest morphology measure for a BENT aggregate, where a
        # cylinder form-factor fit conflates contour length with end-to-end.
        c0 = coords - coords.mean(axis=0)
        rg_list.append(float(np.sqrt((c0 ** 2).sum(axis=1).mean())))
        S = (c0.T @ c0) / len(c0)
        lam = np.sort(np.linalg.eigvalsh(S))          # L1 <= L2 <= L3
        Rg2 = lam.sum()
        b_asph = lam[2] - 0.5 * (lam[0] + lam[1])
        c_acyl = lam[1] - lam[0]
        kappa2 = (b_asph ** 2 + 0.75 * c_acyl ** 2) / Rg2 ** 2
        shape_list.append((lam[0], lam[1], lam[2], b_asph, c_acyl, kappa2))

        # --- ring / closed-loop test (TOPOLOGICAL, not shape-based)
        if a.ring_test:
            # principal plane = the two LARGEST-eigenvalue directions; the ring
            # normal is the smallest. eigh returns ascending eigenvalues.
            _, evec = np.linalg.eigh(S)
            e_small, e_mid, e_big = evec[:, 0], evec[:, 1], evec[:, 2]
            x = c0 @ e_big
            y = c0 @ e_mid
            rho = np.hypot(x, y)                      # in-plane radius
            th = np.arctan2(y, x)
            NB_ANG = 36                               # 10-degree bins
            cnt, _ = np.histogram(th, bins=NB_ANG, range=(-np.pi, np.pi))
            occupied = cnt >= 0.15 * cnt.mean()       # bin holds real mass
            coverage = occupied.sum() / NB_ANG        # 1.0 = closed loop
            # hollowness: a torus has (almost) nothing near the centroid.
            R_char = np.median(rho)
            core_fill = float(np.mean(rho < 0.4 * R_char))
            # A ring is a WORM joined end to end, so the aggregate must be
            # elongated in the first place. Without this gate a compact globule
            # trips the test: any roughly round object has mass at all angles,
            # and a slightly depleted centre then satisfies core_fill too.
            # (Seen on the 25 mM spherical micelles: kappa^2 = 0.04-0.09 yet
            # 12-15% of frames were being called rings.)
            elongated = (kappa2 >= 0.15) and (np.sqrt(lam[2] / lam[0]) >= 2.0)
            is_ring = elongated and (coverage >= 0.95) and (core_fill < 0.08)
            ring_list.append((t, coverage, core_fill, R_char,
                              1 if is_ring else 0, 1 if elongated else 0))

        # --- intra-aggregate distance histogram (no PBC: already unwrapped)
        d = distance_array(coords.astype(np.float32), coords.astype(np.float32))
        iu = np.triu_indices(N, k=1)
        dv = d[iu]
        rmax_frame = float(dv.max())
        rmax_seen = max(rmax_seen, rmax_frame)

        if hist_acc is None:
            # size the histogram generously; grown if a later frame is larger
            nbins = int(np.ceil((rmax_frame * 1.5) / a.dr)) + 1
            hist_acc = np.zeros(nbins)
        need = int(np.ceil(rmax_frame / a.dr)) + 1
        if need > nbins:
            grown = np.zeros(need)
            grown[:nbins] = hist_acc
            hist_acc, nbins = grown, need

        idx = (dv / a.dr).astype(np.int64)
        hist_acc += np.bincount(idx, minlength=nbins)[:nbins]

        nagg_list.append(len(big))
        nbead_list.append(N)
        n_used += 1

    if n_used == 0:
        sys.exit(f"no usable frames (percolating: {n_perc}, too small: {n_small}). "
                 f"A box-spanning aggregate has no well-defined P(q). "
                 f"Use --local for a percolating-safe cross-section.")

    # -----------------------------------------------------------------------
    # --axis: real-space cross-section (topology-independent).
    # -----------------------------------------------------------------------
    if a.axis:
        if not axis_list:
            sys.exit("no frames yielded a usable local-axis estimate")
        A = np.array(axis_list)
        R_rms, R_head, ani = A[:, 0], A[:, 1], A[:, 2]
        print(f"frames used: {n_used}   (--axis: real-space local cylinder axis, "
              f"neighbourhood {a.axis_nb} nm)")
        print(f"<N_agg {a.surf}> = {np.mean(nagg_list):.1f}   <N_beads> = {np.mean(nbead_list):.0f}")
        print("\n--- cross-section from local axes (works for ANY topology) ---")
        print(f"  R from <d^2> (uniform-cylinder) : {R_rms.mean():.3f} +/- {R_rms.std():.3f} nm")
        if np.isfinite(R_head).any():
            print(f"  <d> of HEADGROUPS (surface)     : {np.nanmean(R_head):.3f} +/- {np.nanstd(R_head):.3f} nm")
            print("     (independent estimate: heads sit at the surface, assumes nothing "
                  "about the density profile)")
        print(f"  neighbourhood anisotropy L3/L1  : {ani.mean():.2f}")
        if ani.mean() < 1.5:
            print("     WARNING: < 1.5 means the neighbourhood is nearly isotropic, so the")
            print("     local axis is poorly defined -- increase --axis-nb, or the object")
            print("     may not be locally cylindrical (e.g. a sphere).")
        np.savetxt("axis_R.dat", A, header="R_from_d2(nm)  R_head(nm)  anisotropy",
                   fmt="%10.4f")
        print("  wrote axis_R.dat")
        sys.exit(0)

    # -----------------------------------------------------------------------
    # --local: report the reciprocal-space S(q) and the cross-sectional radius.
    # -----------------------------------------------------------------------
    if a.local:
        qc = 0.5 * (qbins[:-1] + qbins[1:])
        S = sq_acc / n_used
        good = (sq_cnt > 0) & (S > 0)
        qc, S = qc[good], S[good]
        print(f"frames used: {n_used}   (--local: exact reciprocal-space S(q), "
              f"percolation-safe)")
        print(f"<N_agg {a.surf}> = {np.mean(nagg_list):.1f}   <N_beads> = {np.mean(nbead_list):.0f}")
        print(f"q grid: {len(qc)} shells up to {a.qmax_local} nm^-1, "
              f"q_min = 2*pi/L = {qc[0]:.3f} nm^-1")
        print("NOTE: --local gives NO Rg / kappa^2 / D_max / contour length -- those are")
        print("      genuinely undefined for a percolating aggregate. Cross-section only.")

        # Cross-sectional Guinier on ln[q S(q)] vs q^2.
        # CAUTION: unlike the Debye P(q), this S(q) contains INTER-aggregate
        # correlations (other aggregates + the periodic images of this one), so
        # S(q) = P(q)*S_inter(q). The contamination is worst at low q and dies
        # out as S_inter -> 1 at high q. Calibrated against the 50 mM system,
        # where the Debye route gives the truth (Rc = 1.663 nm), a single
        # low-q window over-reads by ~13%. So DO NOT report one window: scan
        # several and quote the median + spread as the real uncertainty.
        print("\n--- cross-sectional Guinier  ln[q S(q)] vs q^2  (LOCAL, percolation-safe) ---")
        if a.guinier_qlo and a.guinier_qhi:
            windows = [(a.guinier_qlo, a.guinier_qhi)]
        else:
            lo0 = max(qc[0] * 2.0, 0.5)
            windows = [(lo0 + 0.1 * k, lo0 + 0.1 * k + 0.40) for k in range(6)]
        res = []
        for lo, hi in windows:
            m = (qc >= lo) & (qc <= hi) & (S > 0)
            if m.sum() < 4:
                continue
            x = qc[m] ** 2
            y = np.log(qc[m] * S[m])
            slope, intercept = np.polyfit(x, y, 1)
            if slope >= 0:
                continue
            Rc_w = float(np.sqrt(-2.0 * slope))
            yhat = slope * x + intercept
            ss_res = float(np.sum((y - yhat) ** 2)); ss_tot = float(np.sum((y - y.mean()) ** 2))
            r2_w = 1.0 - ss_res / ss_tot if ss_tot > 0 else np.nan
            res.append((lo, hi, Rc_w, r2_w))
            print(f"    q = {lo:.2f}-{hi:.2f}:  Rc = {Rc_w:.3f}  R = {np.sqrt(2)*Rc_w:.3f} nm"
                  f"   R^2 = {r2_w:.4f}")
        if res:
            # Keep only the best-linearity windows: curvature in ln[qS] vs q^2 IS
            # the structure-factor contamination, so high R^2 selects the windows
            # where S_inter is flattest. (Validated on 50 mM: this recovers
            # R = 2.40 vs the Debye truth 2.351, ~2%; a median over ALL windows
            # is dragged ~5% low by the contaminated low-q ones.)
            r2s = np.array([r[3] for r in res])
            keep = [r for r, s in zip(res, r2s) if s >= r2s.max() - 0.03]
            arr = np.array([r[2] for r in keep])
            fit_r2 = float(np.max(r2s))
            Rc = float(np.median(arr))
            spread = float(arr.max() - arr.min()) if len(arr) > 1 else 0.0
            print(f"  --> using {len(keep)}/{len(res)} windows with R^2 >= {r2s.max()-0.03:.4f} "
                  f"(flattest S_inter)")
            print(f"  --> Rc = {Rc:.3f} nm (median, spread {spread:.3f})")
            print(f"  --> R  = sqrt(2)*Rc = {np.sqrt(2)*Rc:.3f} nm "
                  f"(window spread {np.sqrt(2)*spread:.3f} -- this is NOT the total "
                  f"uncertainty, see below)")
            print("  NB --local carries ~10-15% systematic uncertainty from inter-aggregate")
            print("     structure (S(q) = P(q)*S_inter). Where the Debye route works "
                  "(non-percolating), PREFER IT.")
        else:
            Rc = fit_r2 = np.nan
            print("  no usable cross-sectional Guinier window.")

        with open(a.out, "w") as fh:
            fh.write(f"# LOCAL reciprocal-space S(q), {a.sim_id}, {n_used} frames\n")
            if np.isfinite(Rc):
                fh.write(f"# cross-sectional Rc = {Rc:.3f} nm, R = {np.sqrt(2)*Rc:.3f} nm, R^2 = {fit_r2:.4f}\n")
            fh.write("# q(nm^-1)      S(q)          q*S(q)\n")
            for qi, Si in zip(qc, S):
                fh.write(f"{qi:12.5f}  {Si:12.6e}  {qi*Si:12.6e}\n")
        print(f"wrote {a.out}")
        sys.exit(0)

    # --- Debye sum from the accumulated histogram
    # P(q) = (1/N^2) [ N + 2 * SUM_{i<j} sinc(q r_ij) ]
    r_edges = (np.arange(nbins) + 0.5) * a.dr
    h = hist_acc / n_used                      # mean pair count per frame per bin
    Nmean = float(np.mean(nbead_list))

    q = np.linspace(a.qmin, a.qmax, a.nq)
    nz = h > 0
    rr = r_edges[nz]
    hh = h[nz]
    P = np.empty_like(q)
    for k, qq in enumerate(q):
        x = qq * rr
        sinc = np.sin(x) / x
        P[k] = (Nmean + 2.0 * np.sum(hh * sinc)) / Nmean ** 2

    with open(a.out, "w") as fh:
        fh.write(f"# single-aggregate P(q) by Debye summation, from {a.sim_id}\n")
        fh.write(f"# frames used {n_used} (stride {a.stride}); percolating skipped {n_perc}; "
                 f"too-small skipped {n_small}\n")
        fh.write(f"# <N_agg {a.surf}> = {np.mean(nagg_list):.1f}   <N_beads> = {Nmean:.0f}"
                 f"{' (incl. bound ' + a.anion + ')' if a.anion else ''}\n")
        fh.write(f"# <Rg> = {np.mean(rg_list):.3f} nm   max intra-aggregate distance = {rmax_seen:.2f} nm\n")
        sh = np.array(shape_list)
        fh.write(f"# gyration tensor: <sqrt(L3/L1)> = {np.mean(np.sqrt(sh[:,2]/sh[:,0])):.2f}"
                 f"  <kappa^2> = {np.mean(sh[:,5]):.3f}"
                 f"  (0 = sphere, 1 = thin rod)\n")
        fh.write(f"# Guinier knee ~ 2*pi/D_max = {2*np.pi/rmax_seen:.3f} nm^-1 (P(q) flattens below this;\n"
                 f"#   this is NOT a truncation floor -- the Debye sum is exact at all q)\n")
        fh.write(f"# expected rod regime (q^-1) if cylindrical: ~{2*np.pi/rmax_seen:.2f} > q > ~1/R\n")
        fh.write("# q(nm^-1)      P(q)          q*P(q)\n")
        for qi, Pi in zip(q, P):
            fh.write(f"{qi:12.5f}  {Pi:12.6e}  {qi*Pi:12.6e}\n")

    print(f"frames used: {n_used}   skipped: {n_perc} percolating, {n_small} too small")
    print(f"<N_agg {a.surf}> = {np.mean(nagg_list):.1f}   <N_beads> = {Nmean:.0f}")
    print(f"<Rg> = {np.mean(rg_list):.3f} nm   D_max = {rmax_seen:.2f} nm")

    # --- model-free shape report
    sh = np.array(shape_list)
    L1, L2, L3 = sh[:, 0], sh[:, 1], sh[:, 2]
    kappa2 = sh[:, 5]
    ax = np.sqrt(L3 / L1)      # principal axis ratio (semi-axis units)
    mid = np.sqrt(L2 / L1)
    print("\n--- gyration tensor (model-free shape) ---")
    print(f"  principal semi-axis ratios  a3:a2:a1 = "
          f"{np.mean(ax):.2f} : {np.mean(mid):.2f} : 1.00")
    print(f"  asphericity  b  = {np.mean(sh[:,3]):.3f} +/- {np.std(sh[:,3]):.3f} nm^2   (0 = sphere)")
    print(f"  acylindricity c = {np.mean(sh[:,4]):.3f} +/- {np.std(sh[:,4]):.3f} nm^2   (0 = cylindrically symmetric)")
    print(f"  rel. shape anisotropy kappa^2 = {np.mean(kappa2):.3f} +/- {np.std(kappa2):.3f}")
    print("     reference: 0 = sphere, 0.25 = flat disc, ->1 = thin rod")
    k = float(np.mean(kappa2))
    if k < 0.05:
        verdict = "essentially SPHERICAL"
    elif k < 0.15:
        verdict = "mildly aspherical (globular/prolate)"
    elif k < 0.4:
        verdict = "clearly PROLATE / short or bent rod"
    else:
        verdict = "strongly ROD-LIKE / elongated worm"
    print(f"  => {verdict}")
    print(f"  NB for a BENT worm kappa^2 under-reads elongation (contour length is\n"
          f"     longer than the end-to-end extent the tensor sees).")
    print(f"Guinier knee ~ 2*pi/D_max = {2*np.pi/rmax_seen:.3f} nm^-1 "
          f"(P(q) flattens below this -- real, not truncation)")
    print(f"  NB the Debye sum is EXACT at all q; unlike the g(r) route it has no "
          f"truncation floor (that was {2*np.pi/(0.5*float(np.min(L_nm))):.3f} nm^-1).")
    if n_perc:
        print(f"NOTE: {n_perc} frames had a box-spanning aggregate (no consistent unwrap); "
              f"P(q) is undefined for those and they were excluded.")
    # -----------------------------------------------------------------------
    # Cross-sectional Guinier.
    # In the rod regime (1/L << q << 1/R) a cylinder obeys
    #     P(q) ~ (pi/q) * P_cs(q),   P_cs(q) ~ P_cs(0) exp(-q^2 Rc^2 / 2)
    # so  ln[q P(q)]  vs  q^2  is LINEAR with slope = -Rc^2/2, where Rc is the
    # cross-sectional radius of GYRATION. For a uniform circular cross-section
    # the geometric radius is R = sqrt(2) * Rc.
    # This is model-free (no form factor assumed) and is the right way to get a
    # radius for a BENT worm, where a cylinder fit conflates contour with
    # end-to-end length.
    # Window: above the Guinier knee (2*pi/D_max) and within the standard
    # cross-sectional Guinier validity q*Rc <~ 1.3. Rc is unknown a priori, so
    # the upper limit is refined iteratively.
    # -----------------------------------------------------------------------
    qlo = a.guinier_qlo if a.guinier_qlo else 1.5 * (2.0 * np.pi / rmax_seen)
    qhi = a.guinier_qhi if a.guinier_qhi else qlo * 2.5
    Rc = fit_r2 = np.nan
    slope = np.nan
    for _ in range(8):
        m = (q >= qlo) & (q <= qhi) & (P > 0)
        if m.sum() < 5:
            break
        x = q[m] ** 2
        y = np.log(q[m] * P[m])
        slope, intercept = np.polyfit(x, y, 1)
        if slope >= 0:
            break
        Rc_new = float(np.sqrt(-2.0 * slope))
        # goodness of fit
        yhat = slope * x + intercept
        ss_res = float(np.sum((y - yhat) ** 2))
        ss_tot = float(np.sum((y - y.mean()) ** 2))
        fit_r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else np.nan
        if a.guinier_qhi is None:
            qhi_new = 1.3 / Rc_new
            if abs(qhi_new - qhi) < 1e-3:
                Rc = Rc_new
                break
            qhi = qhi_new
        Rc = Rc_new
        if a.guinier_qhi is not None:
            break

    if a.ring_test and ring_list:
        rl = np.array(ring_list)
        t_r, cov, core, Rr, isr = rl[:, 0], rl[:, 1], rl[:, 2], rl[:, 3], rl[:, 4]
        elong = rl[:, 5]
        frac = isr.mean()
        print("\n--- ring / closed-loop test (topological) ---")
        print(f"  angular coverage about centroid : {cov.mean():.3f} +/- {cov.std():.3f}"
              f"   (1.00 = mass at all angles = closed loop)")
        print(f"  core fill (hollowness)          : {core.mean():.3f} +/- {core.std():.3f}"
              f"   (~0 = hollow centre = torus)")
        print(f"  in-plane radius R_ring          : {Rr.mean():.2f} nm")
        print(f"  frames ELONGATED enough to ring : {int(elong.sum())}/{len(elong)}"
              f" = {100*elong.mean():.0f}%   (gate: kappa^2>=0.15 and a3/a1>=2)")
        if elong.mean() < 0.5:
            print("     -> mostly GLOBULAR, so a ring is not geometrically possible here;")
            print("        the gate, not the coverage/hollowness, is what rules it out.")
        print(f"  frames classified RING          : {int(isr.sum())}/{len(isr)} = {100*frac:.0f}%")
        if frac > 0:
            first = t_r[isr > 0].min()
            print(f"  first ring frame at t = {first:.0f} ns")
            # contiguity: is it a persistent closure or intermittent?
            late = isr[t_r >= first]
            print(f"  after first closure: {100*late.mean():.0f}% of frames are rings"
                  f" ({'persistent' if late.mean() > 0.8 else 'INTERMITTENT - opens and recloses'})")
        # Thresholds: the per-frame classifier is noisy (an open worm momentarily
        # bent into a C can satisfy coverage+hollowness), so an occasional hit is
        # NOT a ring. Validated on the open 50 mM control, which scored 10-18%.
        # PERSISTENCE is the real signal -- require a clear majority of frames.
        if frac > 0.9:
            verdict = "CLOSED RING (torus) throughout"
        elif frac > 0.5:
            verdict = "ring forms part-way through the run"
        elif frac > 0.0:
            verdict = ("OPEN aggregate -- NOT a closed ring "
                       f"({100*frac:.0f}% isolated hits = the known false-positive rate)")
        else:
            verdict = "OPEN aggregate (worm/hairpin) -- NOT a closed ring"
        print(f"  => {verdict}")
        np.savetxt("ring_test.dat", rl,
                   header="t(ns)  ang_coverage  core_fill  R_ring(nm)  is_ring  elongated",
                   fmt="%10.2f %10.4f %10.4f %10.3f %3d %3d")
        print("  wrote ring_test.dat")

    print("\n--- cross-sectional Guinier  ln[q P(q)] vs q^2 ---")
    if np.isfinite(Rc):
        print(f"  fit window: q = {qlo:.3f} - {qhi:.3f} nm^-1   "
              f"(q*Rc = {qlo*Rc:.2f} - {qhi*Rc:.2f}; valid if <~1.3)")
        print(f"  cross-sectional Rg:  Rc = {Rc:.3f} nm")
        print(f"  geometric radius:    R  = sqrt(2)*Rc = {np.sqrt(2)*Rc:.3f} nm"
              f"   (uniform circular cross-section)")
        print(f"  linearity R^2 = {fit_r2:.4f}"
              f"{'  -- good, rod regime is real' if fit_r2 > 0.99 else '  -- POOR: rod regime may not exist here'}")
    else:
        print("  could not fit a cross-sectional Guinier region "
              "(no positive-slope-free window / too few points).")
        print("  This is expected if the aggregate is not rod-like.")

    # -----------------------------------------------------------------------
    # p(r): pair distance distribution function.
    # This is the same quantity SANS experimentalists recover by indirect
    # Fourier transform of I(q) -- here it is obtained DIRECTLY (it is the
    # histogram the Debye sum is built from), with no inversion and no
    # regularisation ambiguity.
    # Shape reading:
    #   sphere      -> symmetric bell peaking near 0.5*D_max
    #   prolate/rod -> peak at ~the cross-section diameter, then a long
    #                  near-linear tail out to D_max
    #   bent worm   -> as rod, but the tail droops early (end-to-end < contour)
    #   dumbbell/2 domains -> a shoulder or second hump
    # p(r) must go to 0 at r=0 and at r=D_max; D_max is read off where it lands.
    # -----------------------------------------------------------------------
    pr = h.copy()                       # mean pair count per frame per bin
    last = np.max(np.nonzero(pr)[0]) if np.any(pr) else 0
    r_pr = r_edges[:last + 1]
    pr = pr[:last + 1]
    area = np.trapezoid(pr, r_pr) if hasattr(np, "trapezoid") else np.trapz(pr, r_pr)
    pr_norm = pr / area if area > 0 else pr

    # moments: Rg can be recomputed from p(r) as a consistency check
    #   Rg^2 = INT r^2 p(r) dr / (2 INT p(r) dr)
    num = np.trapezoid(pr * r_pr ** 2, r_pr) if hasattr(np, "trapezoid") else np.trapz(pr * r_pr ** 2, r_pr)
    rg_from_pr = float(np.sqrt(num / (2.0 * area))) if area > 0 else float("nan")
    r_peak = float(r_pr[int(np.argmax(pr))])

    with open(a.pr_out, "w") as fh:
        fh.write(f"# p(r) pair distance distribution, {a.sim_id}  ({n_used} frames)\n")
        fh.write(f"# <N_agg {a.surf}> = {np.mean(nagg_list):.1f}   <N_beads> = {Nmean:.0f}\n")
        fh.write(f"# D_max = {r_pr[-1]:.2f} nm   peak at r = {r_peak:.2f} nm "
                 f"({r_peak/r_pr[-1]:.2f} of D_max; ~0.5 for a sphere, much less for a rod)\n")
        fh.write(f"# Rg from p(r) = {rg_from_pr:.3f} nm  (direct tensor value {np.mean(rg_list):.3f} nm)\n")
        fh.write("# r(nm)        p(r) norm      pairs/frame\n")
        for ri, pn, pv in zip(r_pr, pr_norm, pr):
            fh.write(f"{ri:12.5f}  {pn:12.6e}  {pv:12.6e}\n")

    if np.isfinite(Rc):
        with open(a.out) as fh:
            body = fh.read()
        with open(a.out, "w") as fh:
            fh.write(f"# cross-sectional Guinier: Rc = {Rc:.3f} nm, R = sqrt(2)Rc = "
                     f"{np.sqrt(2)*Rc:.3f} nm, fit q = {qlo:.3f}-{qhi:.3f} nm^-1, R^2 = {fit_r2:.4f}\n")
            fh.write(body)

    print(f"wrote {a.out}   (col 1 = P(q), col 2 = q*P(q) Holtzer)")
    print(f"wrote {a.pr_out}  (col 1 = p(r) normalised, col 2 = pairs/frame)")
    print(f"  p(r): D_max = {r_pr[-1]:.2f} nm, peak at r = {r_peak:.2f} nm "
          f"= {r_peak/r_pr[-1]:.2f} D_max  (~0.5 sphere, << 0.5 rod)")
    print(f"  Rg from p(r) = {rg_from_pr:.3f} nm vs {np.mean(rg_list):.3f} nm direct "
          f"({'consistent' if abs(rg_from_pr-np.mean(rg_list)) < 0.15*np.mean(rg_list) else 'MISMATCH - check'})")
