#!/usr/bin/env python3
"""
Counterion binding fraction beta for a micellar CG trajectory.

beta = (bound counterions) / (aggregated surfactant), using a 10 Angstrom
contact criterion (same cutoff/PBC convention as cmc.py).

  - Surfactant (e.g. TTA) is clustered by 10 A contact (PBC-aware, union-find,
    identical to cmc.py). Clusters with >= --min-size molecules are "micelles".
  - A counterion (e.g. ABZO) is "bound" if ANY of its beads is within --cutoff
    of ANY surfactant bead belonging to a micelle.
  - beta(frame) = (# bound counterions) / (# surfactant molecules in micelles).
    Reported as the time average over [--begin,--end] with a std over frames.

NB this is the SIMULATION contact-beta; it is NOT the same quantity as the
experimental conductivity beta=1-alpha. Compare trends, not absolute values.
See memory: project-experimental-cmc.

Loads {cwd}-run-vis.tpr (fallback .gro) + {sim_id}-run-vis.xtc, same as cmc.py.

usage: beta.py SIM_ID SURF ANION [--cutoff 10] [--min-size 20]
                [--begin NS] [--end NS] [--stride N]
"""
import os
import sys
import argparse
from collections import defaultdict

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


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


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Counterion binding fraction beta (10 A contact)")
    p.add_argument("sim_id")
    p.add_argument("surf",  help="surfactant resname (e.g. TTA)")
    p.add_argument("anion", help="counterion resname (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 surfactant count for a cluster to be a micelle")
    p.add_argument("--begin", type=float, default=None, help="start time, ns (default: whole traj)")
    p.add_argument("--end",   type=float, default=None, help="end time, ns")
    p.add_argument("--stride", type=int, default=1, help="frame stride")
    p.add_argument("--by-cluster", action="store_true",
                   help="also report beta resolved PER CLUSTER (largest vs rest, and vs size). "
                        "Use when large and small aggregates coexist: binding is "
                        "curvature-dependent, and the pooled beta averages that away.")
    p.add_argument("--big-cut", type=int, default=150,
                   help="with --by-cluster: clusters >= this are 'large phase' (default 150)")
    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}")
    anion = u.select_atoms(f"resname {a.anion}")
    if len(surf) == 0 or len(anion) == 0:
        sys.exit(f"empty selection: surf={len(surf)} anion={len(anion)}")

    surf_res  = sorted({at.resindex for at in surf.atoms})
    surf_a2r  = np.array([at.resindex for at in surf.atoms])
    anion_a2r = np.array([at.resindex for at in anion.atoms])
    n_anion_total = len({at.resindex for at in anion.atoms})

    betas, nagg_mean, nbound = [], [], []
    percluster = []          # (cluster_size, beta_of_that_cluster)
    times = []
    n_used = 0
    for ts in u.trajectory[::a.stride]:
        t = ts.time / 1000.0  # ps -> ns
        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

        # cluster surfactant by contact
        pairs = capped_distance(surf.positions, surf.positions, a.cutoff,
                                box=box, return_distances=False)
        neigh = set()
        for ai, aj in pairs:
            ri, rj = int(surf_a2r[ai]), int(surf_a2r[aj])
            if ri != rj:
                neigh.add((ri, rj))
        clusters = union_find(surf_res, neigh)
        micelle_res = set()
        for c in clusters:
            if len(c) >= a.min_size:
                micelle_res.update(c)
        n_mic_surf = len(micelle_res)
        if n_mic_surf == 0:
            continue

        # surfactant beads that belong to a micelle
        mic_mask = np.array([r in micelle_res for r in surf_a2r])
        mic_surf_pos = surf.positions[mic_mask]

        # anion beads within cutoff of any micelle surfactant bead
        cp = capped_distance(anion.positions, mic_surf_pos, a.cutoff,
                             box=box, return_distances=False)
        bound_anion_res = {int(anion_a2r[ai]) for ai, _ in cp}
        nb = len(bound_anion_res)

        betas.append(nb / n_mic_surf)
        nagg_mean.append(n_mic_surf)
        nbound.append(nb)
        times.append(t)
        n_used += 1

        # --- per-cluster beta (curvature / phase resolved)
        if a.by_cluster:
            for c in clusters:
                if len(c) < a.min_size:
                    continue
                cset = set(c)
                cmask = np.array([r in cset for r in surf_a2r])
                cpos = surf.positions[cmask]
                cp2 = capped_distance(anion.positions, cpos, a.cutoff,
                                      box=box, return_distances=False)
                nb_c = len({int(anion_a2r[ai]) for ai, _ in cp2})
                percluster.append((len(c), nb_c / len(c)))

    if n_used == 0:
        sys.exit("no frames in the selected window")

    betas = np.array(betas)
    out = f"beta_{a.anion}.dat"
    with open(out, "w") as fh:
        fh.write(f"# beta = bound {a.anion} / aggregated {a.surf}  (cutoff {a.cutoff} A, min-size {a.min_size})\n")
        fh.write(f"# t(ns)   beta   n_aggregated_surf   n_bound_anion\n")
        for t, b, na, nb in zip(times, betas, nagg_mean, nbound):
            fh.write(f"{t:10.2f}  {b:8.4f}  {na:6d}  {nb:6d}\n")

    print(f"frames used: {n_used}  (window "
          f"{a.begin if a.begin is not None else times[0]:.0f}"
          f"-{a.end if a.end is not None else times[-1]:.0f} ns)")
    print(f"aggregated {a.surf}: mean {np.mean(nagg_mean):.0f} / {len(surf_res)} total")
    print(f"bound {a.anion}:     mean {np.mean(nbound):.0f} / {n_anion_total} total")
    print(f"beta = bound_{a.anion} / aggregated_{a.surf} = "
          f"{betas.mean():.3f} +/- {betas.std():.3f}")
    print(f"wrote {out}")

    if a.by_cluster and percluster:
        pc = np.array(percluster)
        sizes, bes = pc[:, 0], pc[:, 1]
        big = sizes >= a.big_cut
        small = ~big
        print(f"\n--- beta resolved by cluster (two-phase check, cut = {a.big_cut}) ---")
        if big.any():
            print(f"  LARGE (>={a.big_cut}): n={big.sum():6d} obs   <size>={sizes[big].mean():6.1f}   "
                  f"beta = {bes[big].mean():.3f} +/- {bes[big].std():.3f}")
        if small.any():
            print(f"  SMALL (<{a.big_cut}):  n={small.sum():6d} obs   <size>={sizes[small].mean():6.1f}   "
                  f"beta = {bes[small].mean():.3f} +/- {bes[small].std():.3f}")
        if big.any() and small.any():
            d = bes[big].mean() - bes[small].mean()
            print(f"  difference (large - small) = {d:+.3f}"
                  f"   [expected POSITIVE: lower curvature -> denser surface charge -> more binding]")
        print("  beta vs size:")
        for lo, hi in [(20, 40), (40, 60), (60, 100), (100, 200), (200, 400), (400, 10000)]:
            m = (sizes >= lo) & (sizes < hi)
            if m.sum() > 10:
                print(f"    size {lo:4d}-{hi:<5d}: n={m.sum():6d}  <size>={sizes[m].mean():6.1f}  "
                      f"beta = {bes[m].mean():.3f} +/- {bes[m].std():.3f}")
        np.savetxt(f"beta_{a.anion}_bycluster.dat", pc,
                   header="cluster_size  beta", fmt="%10.4f")
        print(f"  wrote beta_{a.anion}_bycluster.dat")
