#!/usr/bin/env python3
"""
Extract a single equilibrated micelle from a Dry Martini frame to use as a
seed for a dissolution / persistence run at a lower concentration.

Loads a structure (GRO or TPR+coords), clusters TTA+anion residues by the
same 10 A inter-residue contact criterion as cmc.py, then writes ONE cluster
(unwrapped across PBC) to a seed .gro.  By default it picks the cluster whose
size is closest to the median aggregate size (>= min_size) — a "typical"
micelle rather than the largest fused one.

Usage:
  python extract_micelle.py STRUCT [--coords XTC] [--frame N]
         [--anion BzO] [--surf TTA] [--target N | --median]
         [--min-size 10] [--thresh 10.0] [-o seed_micelle.gro]

Examples:
  python extract_micelle.py 5.0mM-run.gro
  python extract_micelle.py 5.0mM-run.gro --target 55 -o seed55.gro
"""
import argparse
from collections import defaultdict, deque
import numpy as np
import MDAnalysis as mda
from MDAnalysis.lib.distances import capped_distance


def union_find(nodes, pairs):
    parent = {n: n for n in nodes}
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra
    groups = defaultdict(list)
    for n in nodes:
        groups[find(n)].append(n)
    return list(groups.values())


def _reorder_by_species(gro_path, surf, anion):
    """Rewrite a GRO so all `surf` residues come first, then all `anion` —
    matching a [molecules] block order of surf then anion. Coordinates travel
    with their atoms (residues are only relisted, not moved in space). Residues
    are split by contiguous atoms up to each species' bead count, robust to
    resid wrap-around. Any residue whose name is neither surf nor anion is
    appended last in original order (shouldn't happen for a 2-species box)."""
    with open(gro_path) as f:
        lines = f.readlines()
    title = lines[0].rstrip("\n")
    n = int(lines[1])
    atoms = lines[2:2 + n]
    box = lines[2 + n]

    # Split into residues at the natural GRO boundary: a change in the
    # resid+resname field (columns 0:10). No bead-count assumptions, so it
    # is robust to consecutive same-name residues and to per-species bead
    # counts. (2000 residues never wrap the 5-digit resid, so adjacent
    # residues always differ in cols 0:10.)
    residues = []
    cur, cur_key = [], None
    for ln in atoms:
        key = ln[0:10]
        if cur_key is None:
            cur_key = key
        if key != cur_key:
            residues.append((cur[0][5:10].strip(), cur))
            cur, cur_key = [], key
        cur.append(ln)
    if cur:
        residues.append((cur[0][5:10].strip(), cur))

    surf_res = [r for r in residues if r[0] == surf]
    an_res = [r for r in residues if r[0] == anion]
    other = [r for r in residues if r[0] not in (surf, anion)]

    out_lines, serial, resid = [], 0, 0
    for rn, block in surf_res + an_res + other:
        resid += 1
        rid = resid % 100000
        for at in block:
            serial += 1
            out_lines.append(f"{rid:>5d}{rn:<5s}{at[10:15]}{serial % 100000:>5d}"
                             f"{at[20:].rstrip(chr(10))}\n")

    with open(gro_path, "w") as f:
        f.write(title + "\n")
        f.write(f"{n}\n")
        f.writelines(out_lines)
        f.write(box if box.endswith("\n") else box + "\n")
    print(f"  reordered residues -> [{len(surf_res)} {surf}][{len(an_res)} {anion}]"
          + (f"[{len(other)} other]" if other else ""))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("struct")
    p.add_argument("--coords", default=None, help="optional trajectory (XTC); else use STRUCT coords")
    p.add_argument("--frame", type=int, default=-1, help="frame index if --coords given (default last)")
    p.add_argument("--anion", default="BZO")
    p.add_argument("--surf", default="TTA")
    p.add_argument("--target", type=int, default=None, help="pick cluster nearest this size")
    p.add_argument("--median", action="store_true", help="pick cluster nearest median size (default)")
    p.add_argument("--min-size", type=int, default=10)
    p.add_argument("--thresh", type=float, default=10.0)
    p.add_argument("-o", "--out", default="seed_micelle.gro")
    # --- optional box-building (--build) ---
    p.add_argument("--build", action="store_true",
                   help="also build the seeded box with gmx insert-molecules")
    p.add_argument("--box", type=float, help="cubic box edge (nm) for the seeded system")
    p.add_argument("--surf-ci", default="../TTA_SC1.gro", help="single-surfactant GRO for backfill")
    p.add_argument("--anion-ci", default="../BZO-em.gro", help="single-anion GRO for backfill")
    p.add_argument("--ntotal", type=int, default=1000, help="target #molecules of each species")
    p.add_argument("--build-out", default=None, help="output seeded GRO (default {box}mM_seeded.gro)")
    p.add_argument("--gmx", default="gmx", help="gmx executable")
    a = p.parse_args()
    if a.build and not a.box:
        raise SystemExit("--build requires --box (cubic edge in nm)")

    u = mda.Universe(a.struct, a.coords) if a.coords else mda.Universe(a.struct)
    if a.coords:
        u.trajectory[a.frame]

    sel = u.select_atoms(f"resname {a.surf} or resname {a.anion}")
    if len(sel) == 0:
        raise SystemExit(f"No atoms matched resname {a.surf} or {a.anion}. "
                         f"Available: {sorted(set(u.residues.resnames))}")
    box = u.dimensions
    pos = sel.positions
    atom_res = np.array([at.resindex for at in sel])
    res_ids = sorted(set(atom_res.tolist()))

    pairs, _ = capped_distance(pos, pos, a.thresh, box=box)
    edges = set()
    for i, j in pairs:
        ri, rj = int(atom_res[i]), int(atom_res[j])
        if ri != rj:
            edges.add((ri, rj))

    clusters = [c for c in union_find(res_ids, edges) if len(c) >= a.min_size]
    if not clusters:
        raise SystemExit(f"No clusters >= {a.min_size} found.")
    # Report each cluster's size and surfactant/anion composition
    def compo(c):
        ns = sum(1 for ri in c if u.residues[ri].resname == a.surf)
        return ns, len(c) - ns
    clusters.sort(key=len)
    print(f"Found {len(clusters)} clusters >= {a.min_size}:")
    for c in clusters:
        ns, na = compo(c)
        print(f"  N={len(c):3d}   {a.surf}={ns:3d}  {a.anion}={na:3d}"
              f"   (anion/surf ratio {na/ns:.2f})" if ns else "")
    sizes = sorted(len(c) for c in clusters)

    # Match target/median on the SURFACTANT (N_agg) count, matching the
    # cmc.py Data-{surf} convention — not the combined TTA+anion count.
    surf_counts = sorted(compo(c)[0] for c in clusters)
    if a.target is not None:
        goal = a.target
    else:
        goal = int(np.median(surf_counts))
    chosen = min(clusters, key=lambda c: abs(compo(c)[0] - goal))
    n_agg_surf = compo(chosen)[0]
    print(f"Target N_agg({a.surf}) ~{goal}; extracting micelle with "
          f"{a.surf}={n_agg_surf}, {a.anion}={compo(chosen)[1]} "
          f"({len(chosen)} residues total)")

    # Gather all atoms of the chosen residues, unwrap across PBC around atom 0
    grp = u.atoms[np.isin(u.atoms.resindices, np.array(chosen))]
    gpos = grp.positions.copy()
    ref = gpos[0]
    L = box[:3]
    d = gpos - ref
    d -= L * np.round(d / L)
    grp.positions = ref + d

    # Count each species for the report
    n_surf, n_an = compo(chosen)

    grp.write(a.out)
    print(f"Wrote {len(grp)} atoms to {a.out}")

    n_surf_fill = a.ntotal - n_surf
    n_an_fill   = a.ntotal - n_an
    if n_surf_fill < 0 or n_an_fill < 0:
        raise SystemExit(f"Seed larger than --ntotal ({a.ntotal}); nothing to backfill.")

    if not a.build:
        print(f"\nNext: center in the target box and backfill monomers, e.g.")
        print(f"  gmx editconf -f {a.out} -c -box B B B -o seed_centered.gro")
        print(f"  gmx insert-molecules -f seed_centered.gro -ci {a.surf_ci} "
              f"-nmol {n_surf_fill} -radius 0.21 -try 500 -o step1.gro")
        print(f"  gmx insert-molecules -f step1.gro -ci {a.anion_ci} "
              f"-nmol {n_an_fill} -radius 0.21 -try 500 -o seeded.gro")
        print(f"\n(or re-run with --build --box <nm> to do this automatically)")
        raise SystemExit(0)

    # -------- build the seeded box automatically --------
    import subprocess, os
    B = a.box
    out = a.build_out or f"{B:g}box_seeded.gro"
    centered = "seed_centered.gro"
    step1 = "seed_step1.gro"

    def run(cmd):
        print("  $ " + " ".join(cmd))
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            print(r.stderr[-2000:])
            raise SystemExit(f"gmx step failed: {' '.join(cmd)}")

    print(f"\nBuilding seeded box ({B} nm), backfilling "
          f"{n_surf_fill} {a.surf} + {n_an_fill} {a.anion} monomers...")
    run([a.gmx, "editconf", "-f", a.out, "-c", "-box", str(B), str(B), str(B),
         "-o", centered])
    run([a.gmx, "insert-molecules", "-f", centered, "-ci", a.surf_ci,
         "-nmol", str(n_surf_fill), "-radius", "0.21", "-try", "500", "-o", step1])
    run([a.gmx, "insert-molecules", "-f", step1, "-ci", a.anion_ci,
         "-nmol", str(n_an_fill), "-radius", "0.21", "-try", "500", "-o", out])

    # --- reorder residues to [all surf][all anion] to match the .top ---
    # The seed cluster's residues come out interleaved (surf+anion mixed) and
    # insert-molecules appends, so the raw file is not grouped. grompp matches
    # atoms positionally against the .top's [molecules] block order, so an
    # interleaved file gives a cascade of atom-name mismatches. Reorder here.
    _reorder_by_species(out, a.surf, a.anion)

    # verify final counts
    v = mda.Universe(out)
    import collections
    counts = collections.Counter(v.residues.resnames)
    print(f"\nFinal box {out}: {dict(counts)}")
    if counts.get(a.surf) != a.ntotal or counts.get(a.anion) != a.ntotal:
        print(f"WARNING: expected {a.ntotal} of each — insert-molecules may not "
              f"have placed all (raise --try or check box size).")
    for f in (centered, step1):
        if os.path.exists(f):
            os.remove(f)
