supp_nulls.py

import bct
import numpy as np
from scipy import io
from tqdm import tqdm

import os

# ===================================================================
# Basic helpers
# ===================================================================

def neighbors_from_dir_adj(A):
    """
    Build outgoing and incoming neighbor sets for a directed adjacency matrix.
    A: (n, n) binary adjacency
    Returns:
        neighbors_out, neighbors_in: lists of sets
    """
    n = A.shape[0]
    neighbors_out = [set() for _ in range(n)]
    neighbors_in  = [set() for _ in range(n)]
    for i in range(n):
        js = np.nonzero(A[i])[0]
        for j in js:
            neighbors_out[i].add(j)
            neighbors_in[j].add(i)
    return neighbors_out, neighbors_in

def apply_swap_directed(A, W, neighbors_out, neighbors_in,
                        a, b, c, d, u1, v1, u2, v2):
    """
    Apply the directed 2-edge swap:
        a->b, c->d   =>   a->d, c->b
    Update A and neighbor sets in-place.
    """
    wab = W[a, b]
    wcd = W[c, d]

    # remove old edges
    A[a, b] = 0
    A[c, d] = 0
    W[a, b] = 0.
    W[c, d] = 0.

    neighbors_out[a].remove(b)
    neighbors_in[b].remove(a)
    neighbors_out[c].remove(d)
    neighbors_in[d].remove(c)

    # add new edges
    A[u1, v1] = 1
    A[u2, v2] = 1
    W[u1, v1] = wab
    W[u2, v2] = wcd

    neighbors_out[u1].add(v1)
    neighbors_in[v1].add(u1)
    neighbors_out[u2].add(v2)
    neighbors_in[v2].add(u2)

def revert_swap_directed(A, W, neighbors_out, neighbors_in,
                         a, b, c, d, u1, v1, u2, v2):
    """
    Revert the directed 2-edge swap.
    """
    wab = W[u1, v1]
    wcd = W[u2, v2]

    # remove new
    A[u1, v1] = 0
    A[u2, v2] = 0
    W[u1, v1] = 0.
    W[u2, v2] = 0.

    neighbors_out[u1].remove(v1)
    neighbors_in[v1].remove(u1)
    neighbors_out[u2].remove(v2)
    neighbors_in[v2].remove(u2)

    # add old
    A[a, b] = 1
    A[c, d] = 1
    W[a, b] = wab
    W[c, d] = wcd

    neighbors_out[a].add(b)
    neighbors_in[b].add(a)
    neighbors_out[c].add(d)
    neighbors_in[d].add(c)


# ===================================================================
# 3–5 cycle counting (directed)
# ===================================================================

def count_3cycles_edge_directed(i, j, neighbors_out, neighbors_in):
    """
    Number of directed 3-cycles involving edge i->j:
        i -> j -> k -> i
    """
    # j -> k and k -> i  => k in out(j) ∩ in(i)
    return len(neighbors_out[j].intersection(neighbors_in[i]))

def count_4cycles_edge_directed(i, j, neighbors_out):
    """
    Number of directed 4-cycles involving edge i->j:
        i -> j -> u -> v -> i
    """
    cnt = 0
    for u in neighbors_out[j]:          # j -> u
        if u == i:
            continue
        for v in neighbors_out[u]:      # u -> v
            if v in (i, j):
                continue
            # closing edge v -> i
            if i in neighbors_out[v]:   # v -> i
                cnt += 1
    return cnt

def count_5cycles_edge_directed(i, j, neighbors_out):
    """
    Number of directed 5-cycles involving edge i->j:
        i -> j -> a -> b -> c -> i
    """
    cnt = 0
    for a in neighbors_out[j]:          # j -> a
        if a == i:
            continue
        for b in neighbors_out[a]:      # a -> b
            if b in (i, j):
                continue
            for c in neighbors_out[b]:  # b -> c
                if c in (i, j, a):
                    continue
                # closing edge c -> i
                if i in neighbors_out[c]:   # c -> i
                    cnt += 1
    return cnt

def global_cycle_counts_directed(A, cycles5=True):
    """
    Compute global counts of 3-, 4-, and 5-cycles.
    """
    neighbors_out, neighbors_in = neighbors_from_dir_adj(A)
    n = A.shape[0]
    C3 = C4 = 0
    C5 = 0 if cycles5 else None
    for i in range(n):
        for j in neighbors_out[i]:
            C3 += count_3cycles_edge_directed(i, j, neighbors_out, neighbors_in)
            C4 += count_4cycles_edge_directed(i, j, neighbors_out)
            if cycles5:
                C5 += count_5cycles_edge_directed(i, j, neighbors_out)
    C3 //= 3
    C4 //= 4
    if cycles5:
        C5 //= 5
    return C3, C4, C5

def energy_cycles(C3, C3_target, C4, C4_target, C5=None, C5_target=None,
                  w3=1.0, w4=1.0, w5=0.0):
    """
    Quadratic energy penalizing deviations from target cycle counts.
    """
    energy = (w3 * (C3 - C3_target) ** 2 +
              w4 * (C4 - C4_target) ** 2)
    if C5 is not None and C5_target is not None and w5 > 0:
        energy += w5 * (C5 - C5_target) ** 2
    return energy

def canon_directed_cycle(seq):
    """
    Canonicalize a directed cycle given as an ordered sequence of nodes.

    Two representations that differ only by rotation are the same cycle:
      (a,b,c,d) == (b,c,d,a) == (c,d,a,b) == (d,a,b,c)

    Reverse direction is NOT considered equivalent (keeps orientation).
    """
    seq = tuple(seq)
    k = len(seq)
    rots = [seq[i:] + seq[:i] for i in range(k)]
    return min(rots)

def count_3cycles_canon_directed(u, v, neighbors_out):
    """
    All distinct directed 3-cycles containing edge u->v.
    Returns set of frozenset({u,v,k}).
    """
    cycles = set()
    for k in neighbors_out[v]:
        if u in neighbors_out[k]:
            cycles.add(canon_directed_cycle((u, v, k)))
    return cycles

def count_4cycles_canon_directed(u, v, neighbors_out):
    """
    All distinct directed 4-cycles containing edge u->v.
    Returns set of frozenset({u,v,x,y}).
    """
    cycles = set()
    for x in neighbors_out[v]:
        if x == u:
            continue
        for y in neighbors_out[x]:
            if y in (u, v):
                continue
            if u in neighbors_out[y]:
                cycles.add(canon_directed_cycle((u, v, x, y)))
    return cycles

def count_5cycles_canon_directed(u, v, neighbors_out):
    """
    All distinct directed 5-cycles containing edge u->v.
    Returns set of frozenset({u,v,x,y,z}).
    """
    cycles = set()
    for x in neighbors_out[v]:
        if x == u:
            continue
        for y in neighbors_out[x]:
            if y in (u, v):
                continue
            for z in neighbors_out[y]:
                if z in (u, v, x):
                    continue
                if u in neighbors_out[z]:
                    cycles.add(canon_directed_cycle((u, v, x, y, z)))
    return cycles

def enumerate_local_cycles(neighbors_out, vertices, k):
    """
    Enumerate all distinct directed k-cycles that contain
    at least one edge with both endpoints in vertices.
    """
    cycles = set()

    # internal edges only
    for u in vertices:
        for v in neighbors_out[u] & vertices:

            if k == 3:
                cycles |= count_3cycles_canon_directed(u, v, neighbors_out)
            elif k == 4:
                cycles |= count_4cycles_canon_directed(u, v, neighbors_out)
            elif k == 5:
                cycles |= count_5cycles_canon_directed(u, v, neighbors_out)
            else:
                raise ValueError("k must be 3, 4, or 5")

    return cycles


# ===================================================================
# Modularity masks (level-1 within vs between)
# ===================================================================

def build_within_between_masks(partition):
    """
    partition: array-like of length n, module label for each node.
    Returns:
        within_mask, between_mask: boolean (n, n) masks.
    """
    partition = np.asarray(partition)
    n = partition.shape[0]
    within = np.zeros((n, n), dtype=bool)
    for m in np.unique(partition):
        idx = np.where(partition == m)[0]
        within[np.ix_(idx, idx)] = True
    between = ~within
    return within, between

def update_edge_sets_on_accept(edges_within, edges_between,
                               within_mask,
                               a, b, c, d, u1, v1, u2, v2,
):
    """
    Update modularity edge sets after accepting a swap.
    """
    # determine category from removed edge
    if within_mask[a, b]:
        edges = edges_within
    else:
        edges = edges_between

    edges.remove((a, b))
    edges.remove((c, d))
    edges.add((u1, v1))
    edges.add((u2, v2))


# ===================================================================
# Degree- and modularity-preserving swap proposal (directed)
# ===================================================================

def propose_modular_swap_directed(A, edges_within, edges_between,
                                  within_mask, between_mask, rs,
                                  max_attempts=40):
    """
    Propose a directed 2-edge swap that:
      - preserves in/out-degree,
      - preserves within/between category of both edges.
      a->b, c->d   =>   a->d (u1->v1), c->b (u2->v2)

    Returns:
        (a,b,c,d,u1,v1,u2,v2,category) or None if no valid swap found.
    """
    # count edges in both categories
    n_within  = len(edges_within)
    n_between = len(edges_between)

    # --- both categories invalid ---
    if n_within < 2 and n_between < 2:
        return None

    # --- must choose "between" ---
    if n_within < 2:
        edges = edges_between
        orig_mask = between_mask

    # --- must choose "within" ---
    elif n_between < 2:
        edges = edges_within
        orig_mask = within_mask

    # --- both categories valid → choose proportionally ---
    else:
        r = rs.random()
        if r < n_within / (n_within + n_between):
            edges = edges_within
            orig_mask = within_mask
        else:
            edges = edges_between
            orig_mask = between_mask

    edges_arr = np.asarray(list(edges), dtype=int)
    m = edges_arr.shape[0]

    for _ in range(max_attempts):

        # sample two distinct edges uniformly
        idx = rs.choice(m, size=2, replace=False)
        (a, b), (c, d) = edges_arr[idx]

        # all distinct nodes
        if len({a, b, c, d}) < 4:
            continue

        # proposed new edges
        u1, v1 = a, d
        u2, v2 = c, b

        # no self loops
        if u1 == v1 or u2 == v2:
            continue

        # no multiedges
        if A[u1, v1] or A[u2, v2]:
            continue

        # must remain in same category
        if not (orig_mask[u1, v1] and orig_mask[u2, v2]):
            continue

        return a, b, c, d, u1, v1, u2, v2

    return None


# ===================================================================
# Modularity + cycles null 
# degree and level-1 modularity -> hard constraints,
# 3–5 cycles -> soft constraint, with modularity-preserving burn-in.
# ===================================================================

def modularity_cycles_null_directed(
    A,
    partition,
    n_burnin=10000,
    n_stage=100,
    n_iter=10000,
    T_init=1000,
    frac=0.5,
    w3=1.0,
    w4=1.0,
    w5=1.0,
    verbose=False,
    seed=None,
):
    """
    Build a null that preserves:
      - in/out-degree exactly,
      - within/between-module edge counts exactly (level-1 modularity),
    and softly preserves (using simulated annealing):
      - global 3-, 4-, 5-cycle counts.

    A: (n, n) directed adjacency
    partition: array-like of length n, level-1 module label for each node.

    Returns
    -------
    W_best: (n, n) null network
    E_best: energy of W_best
    C3_target, C4_target, C5_target: cycle counts in original graph
    Cmean_target: average clustering in original graph
    rand_W: rewired modularity-preserving network (before SA)
    """
    try:
        A = np.asarray(A)
    except TypeError as err:
        msg = ('A must be array_like. Received: {}.'.format(type(A)))
        raise TypeError(msg) from err

    if frac > 1 or frac <= 0:
        msg = ('frac must be between 0 and 1. '
               'Received: {}.'.format(frac))
        raise ValueError(msg)

    rs = bct.get_rng(seed)

    W = A.copy().astype(float)
    A = (W != 0).astype(int).copy()
    np.fill_diagonal(W, 0.)
    np.fill_diagonal(A, 0)

    n = A.shape[0]
    m = A.sum()
    if m < 2 or m > n * (n - 1) - 2:
        raise ValueError("Trivial graph: no degree-preserving nulls possible")

    within_mask, between_mask = build_within_between_masks(partition)
    edges_within  = set(zip(*np.nonzero(A & within_mask)))
    edges_between = set(zip(*np.nonzero(A & between_mask)))

    neighbors_out, neighbors_in = neighbors_from_dir_adj(A)

    cycles5 = True if w5 > 0 else False

    # Target cycle counts from original graph
    C3_target, C4_target, C5_target = global_cycle_counts_directed(A, cycles5=cycles5)

    # Burn-in: modularity- & degree-preserving MS swaps
    for _ in range(n_burnin):
        prop = propose_modular_swap_directed(A, edges_within, edges_between, 
                                             within_mask, between_mask, rs)
        if prop is None:
            continue
        a, b, c, d, u1, v1, u2, v2 = prop
        apply_swap_directed(A, W, neighbors_out, neighbors_in,
                            a, b, c, d, u1, v1, u2, v2)
        update_edge_sets_on_accept(edges_within, edges_between,
                                   within_mask,
                                   a, b, c, d, u1, v1, u2, v2)
    
    rand_W = W.copy()

    # Compute cycle counts
    C3, C4, C5 = global_cycle_counts_directed(A, cycles5=cycles5)
    E = energy_cycles(C3, C3_target, C4, C4_target, C5=C5, C5_target=C5_target,
                      w3=w3, w4=w4, w5=w5)
    E_best = E
    W_best = W.copy()

    # SA phase

    for stage in tqdm(range(n_stage), desc='annealing progress'):
        n_accept = 0
        T = T_init * (frac ** stage)
        for it in range(n_iter):
            prop = propose_modular_swap_directed(A, edges_within, edges_between,
                                                 within_mask, between_mask, rs)
            if prop is None:
                continue
            a, b, c, d, u1, v1, u2, v2 = prop

            vertices = {a, b, c, d}

            # Local cycle counts BEFORE
            C3_old = enumerate_local_cycles(neighbors_out, vertices, k=3)
            C4_old = enumerate_local_cycles(neighbors_out, vertices, k=4)
            if cycles5:
                C5_old = enumerate_local_cycles(neighbors_out, vertices, k=5)

            # Apply swap
            apply_swap_directed(A, W, neighbors_out, neighbors_in,
                                a, b, c, d, u1, v1, u2, v2)

            # Local cycle counts AFTER
            C3_new = enumerate_local_cycles(neighbors_out, vertices, k=3)
            C4_new = enumerate_local_cycles(neighbors_out, vertices, k=4)
            if cycles5:
                C5_new = enumerate_local_cycles(neighbors_out, vertices, k=5)

            dC3 = len(C3_new) - len(C3_old)
            dC4 = len(C4_new) - len(C4_old)
            if cycles5:
                dC5 = len(C5_new) - len(C5_old)

            C3_candidate = C3 + dC3
            C4_candidate = C4 + dC4
            if cycles5:
                C5_candidate = C5 + dC5
            else: C5_candidate = None

            E_candidate = energy_cycles(
                C3_candidate, C3_target,
                C4_candidate, C4_target,
                C5=C5_candidate, C5_target=C5_target,
                w3=w3, w4=w4, w5=w5
            )
            dE = E_candidate - E

            if dE <= 0 or rs.random() < np.exp(-dE / T):
                # Accept
                C3, C4, C5 = C3_candidate, C4_candidate, C5_candidate
                E = E_candidate
                update_edge_sets_on_accept(edges_within, edges_between,
                                           within_mask,
                                           a, b, c, d, u1, v1, u2, v2)
                if E < E_best:
                    E_best = E
                    W_best = W.copy()

                n_accept += 1
            else:
                # Reject
                revert_swap_directed(A, W, neighbors_out, neighbors_in,
                                     a, b, c, d, u1, v1, u2, v2)
        
        if verbose:
            print('\nstage {:d}, temp {:.5f}, best energy {:.5f}, '
                  'frac of accepted moves {:.3f}'.format(stage, T, E_best, 
                                                         n_accept/n_iter))
        if E_best == 0:
            break

    return W_best, E_best, C3_target, C4_target, C5_target, rand_W


# ===================================================================
# Clustering tracking
# ===================================================================

def update_affected(i, affected, Cnode, C_old):
    """
    Track nodes whose clustering coefficient is affected by a swap.
    """
    if i not in affected:
        C_old[i] = Cnode[i]
        affected.add(i)

def update_reciprocal(r, A, u, v, delta,
                      affected, Cnode, C_old,
                      r_old):
    """
    delta = +1 if edge u->v is ADDED
          = -1 if edge u->v is REMOVED
    """
    if A[v, u]: # reciprocal partner exists
        update_affected(u, affected, Cnode, C_old)
        update_affected(v, affected, Cnode, C_old)
        if u not in r_old:
            r_old[u] = r[u]
        if v not in r_old:
            r_old[v] = r[v]
        r[u] += delta
        r[v] += delta

def update_triangles(t, A, neighbors_out, neighbors_in, u, v, delta,
                     affected, Cnode, C_old, t_old):
    """
    Incrementally update weighted triangle counts for an edge toggle u->v.

    delta = +1 if edge u->v was ADDED
          = -1 if edge u->v was REMOVED
    """
    Nu = neighbors_out[u] | neighbors_in[u]
    Nv = neighbors_out[v] | neighbors_in[v]
    common = Nu & Nv

    for w in common:

        # weight of the other two sides of the triangle
        w_other = ((A[v, w] + A[w, v]) *
                   (A[w, u] + A[u, w]))

        if w_other == 0:
            continue

        d = delta * w_other

        update_affected(u, affected, Cnode, C_old)
        update_affected(v, affected, Cnode, C_old)
        update_affected(w, affected, Cnode, C_old)
        if u not in t_old:
            t_old[u] = t[u]
        if v not in t_old:
            t_old[v] = t[v]
        if w not in t_old:
            t_old[w] = t[w]
        t[u] += d
        t[v] += d
        t[w] += d

def clustering_from_parts(tu, ku, ru):
    """
    Compute clustering from components
    """
    denom = ku * (ku - 1) - 2 * ru
    if denom == 0 or tu == 0:
        return 0.0
    return tu / denom

def init_fagiolo_state(A):
    """
    Initialize triangle counts, degree sums, reciprocal degrees, 
    and clustering coefficients
    """
    n = A.shape[0]
    
    S = A + A.T
    t = np.diag(np.dot(S, np.dot(S, S))) / 2
    k = np.sum(S, axis=1)
    r = np.diag(np.dot(A, A)).copy()

    C = np.zeros(n)
    for u in range(n):
        C[u] = clustering_from_parts(t[u], k[u], r[u])

    return t, k, r, C


# ===================================================================
# Modularity + clustering null 
# degree and level-1 modularity -> hard constraints,
# clustering -> soft constraint, with modularity-preserving burn-in.
# ===================================================================

def modularity_clustering_null_directed(
    A,
    partition,
    n_burnin=10000,
    n_stage=100,
    n_iter=10000,
    T_init=1e-3,
    frac=0.5,
    verbose=False,
    seed=None,
):
    """
    Build a null that preserves:
      - in/out-degree exactly,
      - within/between-module edge counts exactly (level-1 modularity),
    and softly preserves (using simulated annealing):
      - average directed clustering coefficient.

    A: (n, n) directed adjacency
    partition: array-like of length n, level-1 module label for each node.

    Returns
    -------
    W_best: (n, n) null network
    E_best: energy of W_best
    Cmean_target: average clustering in original graph
    rand_W: rewired modularity-preserving network (before SA)
    """
    try:
        A = np.asarray(A)
    except TypeError as err:
        msg = ('A must be array_like. Received: {}.'.format(type(A)))
        raise TypeError(msg) from err

    if frac > 1 or frac <= 0:
        msg = ('frac must be between 0 and 1. '
               'Received: {}.'.format(frac))
        raise ValueError(msg)

    rs = bct.get_rng(seed)

    W = A.copy().astype(float)
    A = (W != 0).astype(int).copy()
    np.fill_diagonal(W, 0.)
    np.fill_diagonal(A, 0)

    n = A.shape[0]
    m = A.sum()
    if m < 2 or m > n * (n - 1) - 2:
        raise ValueError("Trivial graph: no degree-preserving nulls possible")

    within_mask, between_mask = build_within_between_masks(partition)
    edges_within  = set(zip(*np.nonzero(A & within_mask)))
    edges_between = set(zip(*np.nonzero(A & between_mask)))

    neighbors_out, neighbors_in = neighbors_from_dir_adj(A)

    # Target clustering from original graph
    Cmean_target = bct.clustering_coef_bd(A.astype(float)).mean()

    # Burn-in: modularity- & degree-preserving MS swaps
    for _ in range(n_burnin):
        prop = propose_modular_swap_directed(A, edges_within, edges_between, 
                                             within_mask, between_mask, rs)
        if prop is None:
            continue
        a, b, c, d, u1, v1, u2, v2 = prop
        apply_swap_directed(A, W, neighbors_out, neighbors_in,
                            a, b, c, d, u1, v1, u2, v2)
        update_edge_sets_on_accept(edges_within, edges_between,
                                   within_mask,
                                   a, b, c, d, u1, v1, u2, v2)
        
    rand_W = W.copy()

    # Compute clustering
    t, k, r, Cnode = init_fagiolo_state(A)
    Csum = Cnode.sum()
    Cmean = Csum / n

    E = (Cmean - Cmean_target) ** 2
    E_best = E
    W_best = W.copy()

    # SA phase

    for stage in tqdm(range(n_stage), desc='annealing progress'):
        n_accept = 0
        T = T_init * (frac ** stage)
        for it in range(n_iter):
            prop = propose_modular_swap_directed(A, edges_within, edges_between,
                                                 within_mask, between_mask, rs)
            if prop is None:
                continue
            a, b, c, d, u1, v1, u2, v2 = prop

            affected = set()
            C_old = {}
            t_old = {}
            r_old = {}
            Csum_old = Csum

            for u, v, sgn in [(a, b, -1), (c, d, -1)]:
                update_reciprocal(r, A, u, v, sgn,
                                  affected, Cnode, C_old, r_old)
                update_triangles(t, A, neighbors_out, neighbors_in, u, v, sgn,
                                 affected, Cnode, C_old, t_old)

            # Apply swap
            apply_swap_directed(A, W, neighbors_out, neighbors_in,
                                a, b, c, d, u1, v1, u2, v2)
            
            for u, v, sgn in [(a, d, +1), (c, b, +1)]:
                update_reciprocal(r, A, u, v, sgn,
                                  affected, Cnode, C_old, r_old)
                update_triangles(t, A, neighbors_out, neighbors_in, u, v, sgn,
                                 affected, Cnode, C_old, t_old)

            for u in affected:
                Cnode[u] = clustering_from_parts(t[u], k[u], r[u])

            Csum += sum(Cnode[u] - C_old[u] for u in affected)
            Cmean_new = Csum / n

            E_candidate = (Cmean_new - Cmean_target) ** 2
            dE = E_candidate - E

            if dE <= 0 or rs.random() < np.exp(-dE / T):
                # Accept
                Cmean = Cmean_new
                E = E_candidate
                update_edge_sets_on_accept(edges_within, edges_between,
                                           within_mask,
                                           a, b, c, d, u1, v1, u2, v2)
                if E < E_best:
                    E_best = E
                    W_best = W.copy()

                n_accept += 1
            else:
                # Reject
                revert_swap_directed(A, W, neighbors_out, neighbors_in,
                                    a, b, c, d, u1, v1, u2, v2)
                for i, val in t_old.items():
                    t[i] = val
                for i, val in r_old.items():
                    r[i] = val
                for i, val in C_old.items():
                    Cnode[i] = val
                Csum = Csum_old
        
        if verbose:
            print('\nstage {:d}, temp {:.5f}, best energy {:.5f}, '
                  'frac of accepted moves {:.3f}'.format(stage, T, E_best, 
                                                         n_accept/n_iter))
        if E_best == 0:
            break

    return W_best, E_best, Cmean_target, rand_W


# ===================================================================
# Modularity + clustering null 
# degree and level-1 modularity -> hard constraints,
# 3–5 cycles and clustering -> soft constraint, 
# with modularity-preserving burn-in.
# ===================================================================

def energy_cycles_clustering(C3, C3_target, s3, C4, C4_target, s4,
                             Cmean, Cmean_target, sC,
                             C5=None, C5_target=None, s5=None,
                             w3=1.0, w4=1.0, w5=0.0, wC=1.0):
    """
    Quadratic energy penalizing deviations from target cycle counts
    and average clustering coefficient.
    """
    e3 = (C3 - C3_target) / s3
    e4 = (C4 - C4_target) / s4
    eC = (Cmean - Cmean_target) / sC
    energy = (w3 * e3 * e3 +
              w4 * e4 * e4 +
              wC * eC * eC)
    if C5 is not None and C5_target is not None and w5 > 0:
        e5 = (C5 - C5_target) / s5
        energy += w5 * e5 * e5
    return energy

def modularity_cycles_clustering_null_directed(
    A,
    partition,
    n_burnin=10000,
    n_stage=100,
    n_iter=10000,
    T_init=1,
    frac=0.5,
    w3=1.0,
    w4=1.0,
    w5=1.0,
    wC=1.0,
    verbose=False,
    seed=None,
):
    """
    Build a null that preserves:
      - in/out-degree exactly,
      - within/between-module edge counts exactly (level-1 modularity),
    and softly preserves (using simulated annealing):
      - global 3-, 4-, 5-cycle counts,
      - average directed clustering coefficient.

    A: (n, n) directed adjacency
    partition: array-like of length n, level-1 module label for each node.

    Returns
    -------
    W_best: (n, n) null network
    E_best: energy of W_best
    C3_target, C4_target, C5_target: cycle counts in original graph
    Cmean_target: average clustering in original graph
    rand_W: rewired modularity-preserving network (before SA)
    """
    try:
        A = np.asarray(A)
    except TypeError as err:
        msg = ('A must be array_like. Received: {}.'.format(type(A)))
        raise TypeError(msg) from err

    if frac > 1 or frac <= 0:
        msg = ('frac must be between 0 and 1. '
               'Received: {}.'.format(frac))
        raise ValueError(msg)

    rs = bct.get_rng(seed)

    W = A.copy().astype(float)
    A = (W != 0).astype(int).copy()
    np.fill_diagonal(W, 0.)
    np.fill_diagonal(A, 0)

    n = A.shape[0]
    m = A.sum()
    if m < 2 or m > n * (n - 1) - 2:
        raise ValueError("Trivial graph: no degree-preserving nulls possible")

    within_mask, between_mask = build_within_between_masks(partition)
    edges_within  = set(zip(*np.nonzero(A & within_mask)))
    edges_between = set(zip(*np.nonzero(A & between_mask)))

    neighbors_out, neighbors_in = neighbors_from_dir_adj(A)

    cycles5 = True if w5 > 0 else False

    # Target cycle counts from original graph
    C3_target, C4_target, C5_target = global_cycle_counts_directed(A, cycles5=cycles5)
    Cmean_target = bct.clustering_coef_bd(A.astype(float)).mean()

    s3 = max(1.0, abs(C3_target))
    s4 = max(1.0, abs(C4_target))
    s5 = max(1.0, abs(C5_target)) if cycles5 else None
    sC = max(1.0, abs(Cmean_target))

    # Burn-in: modularity- & degree-preserving MS swaps
    for _ in range(n_burnin):
        prop = propose_modular_swap_directed(A, edges_within, edges_between, 
                                             within_mask, between_mask, rs)
        if prop is None:
            continue
        a, b, c, d, u1, v1, u2, v2 = prop
        apply_swap_directed(A, W, neighbors_out, neighbors_in,
                            a, b, c, d, u1, v1, u2, v2)
        update_edge_sets_on_accept(edges_within, edges_between,
                                   within_mask,
                                   a, b, c, d, u1, v1, u2, v2)

    rand_W = W.copy()

    # Compute cycle counts and clustering
    C3, C4, C5 = global_cycle_counts_directed(A, cycles5=cycles5)
    t, k, r, Cnode = init_fagiolo_state(A)
    Csum = Cnode.sum()
    Cmean = Csum / n

    E = energy_cycles_clustering(C3, C3_target, s3, 
                                 C4, C4_target, s4,
                                 Cmean, Cmean_target, sC,
                                 C5=C5, C5_target=C5_target, s5=s5,
                                 w3=w3, w4=w4, w5=w5, wC=wC)
    E_best = E
    W_best = W.copy()

    # SA phase

    for stage in tqdm(range(n_stage), desc='annealing progress'):
        n_accept = 0
        T = T_init * (frac ** stage)
        for it in range(n_iter):
            prop = propose_modular_swap_directed(A, edges_within, edges_between,
                                                 within_mask, between_mask, rs)
            if prop is None:
                continue
            a, b, c, d, u1, v1, u2, v2 = prop

            vertices = {a, b, c, d}

            # Local cycle counts BEFORE
            C3_old = enumerate_local_cycles(neighbors_out, vertices, k=3)
            C4_old = enumerate_local_cycles(neighbors_out, vertices, k=4)
            if cycles5:
                C5_old = enumerate_local_cycles(neighbors_out, vertices, k=5)

            affected = set()
            C_old = {}
            t_old = {}
            r_old = {}
            Csum_old = Csum

            for u, v, sgn in [(a, b, -1), (c, d, -1)]:
                update_reciprocal(r, A, u, v, sgn,
                                  affected, Cnode, C_old, r_old)
                update_triangles(t, A, neighbors_out, neighbors_in, u, v, sgn,
                                 affected, Cnode, C_old, t_old)

            # Apply swap
            apply_swap_directed(A, W, neighbors_out, neighbors_in,
                                a, b, c, d, u1, v1, u2, v2)

            # Local cycle counts AFTER
            C3_new = enumerate_local_cycles(neighbors_out, vertices, k=3)
            C4_new = enumerate_local_cycles(neighbors_out, vertices, k=4)
            if cycles5:
                C5_new = enumerate_local_cycles(neighbors_out, vertices, k=5)

            dC3 = len(C3_new) - len(C3_old)
            dC4 = len(C4_new) - len(C4_old)
            if cycles5:
                dC5 = len(C5_new) - len(C5_old)

            C3_candidate = C3 + dC3
            C4_candidate = C4 + dC4
            if cycles5:
                C5_candidate = C5 + dC5
            else: C5_candidate = None

            for u, v, sgn in [(a, d, +1), (c, b, +1)]:
                update_reciprocal(r, A, u, v, sgn,
                                  affected, Cnode, C_old, r_old)
                update_triangles(t, A, neighbors_out, neighbors_in, u, v, sgn,
                                 affected, Cnode, C_old, t_old)

            for u in affected:
                Cnode[u] = clustering_from_parts(t[u], k[u], r[u])

            Csum += sum(Cnode[u] - C_old[u] for u in affected)
            Cmean_new = Csum / n

            E_candidate = energy_cycles_clustering(
                C3_candidate, C3_target, s3,
                C4_candidate, C4_target, s4,
                Cmean_new, Cmean_target, sC,
                C5=C5_candidate, C5_target=C5_target, s5=s5,
                w3=w3, w4=w4, w5=w5, wC=wC
            )
            dE = E_candidate - E

            if dE <= 0 or rs.random() < np.exp(-dE / T):
                # Accept
                C3, C4, C5 = C3_candidate, C4_candidate, C5_candidate
                Cmean = Cmean_new
                E = E_candidate
                update_edge_sets_on_accept(edges_within, edges_between,
                                           within_mask,
                                           a, b, c, d, u1, v1, u2, v2)
                if E < E_best:
                    E_best = E
                    W_best = W.copy()

                n_accept += 1
            else:
                # Reject
                revert_swap_directed(A, W, neighbors_out, neighbors_in,
                                     a, b, c, d, u1, v1, u2, v2)
                for i, val in t_old.items():
                    t[i] = val
                for i, val in r_old.items():
                    r[i] = val
                for i, val in C_old.items():
                    Cnode[i] = val
                Csum = Csum_old
        
        if verbose:
            print('\nstage {:d}, temp {:.5f}, best energy {:.5f}, '
                  'frac of accepted moves {:.3f}'.format(stage, T, E_best, 
                                                         n_accept/n_iter))
        if E_best == 0:
            break

    return W_best, E_best, C3_target, C4_target, C5_target, Cmean_target, rand_W


# ===================================================================
# Targeted degree-preserving rewiring to destroy short cycles
# ===================================================================

def creates_short_cycle(u, v, neighbors_out, max_cycle_len):
    """
    Would adding edge u->v create a directed cycle
    of length <= max_cycle_len.

    Checks whether there exists a path v -> ... -> u
    of length <= max_cycle_len - 1 via DFS from v.
    """
    max_depth = max_cycle_len - 1
    stack = [(v, 0)]
    visited = {v}

    while stack:
        node, depth = stack.pop()
        if depth >= max_depth:
            continue

        for nxt in neighbors_out[node]:
            if nxt == u:
                return True
            if nxt not in visited:
                visited.add(nxt)
                stack.append((nxt, depth + 1))

    return False

def creates_triangle(i, j, neighbors_out, neighbors_in):
    return bool(neighbors_out[j] & neighbors_in[i])

def propose_targeted_ms_swap_directed(A, neighbors_out, neighbors_in, rs,
                                      max_cycle_len, max_d_tries=300,
                                      max_c_tries=12):
    """
    Targeted degree-preserving MS swap that avoids creating short cycles.
    """

    n = A.shape[0]
    i_list, j_list = np.nonzero(A)
    k = len(i_list)
    if k < 2:
        return None

    # pick first edge a->b
    idx = rs.randint(k)
    a, b = int(i_list[idx]), int(j_list[idx])

    # try to find a good d such that a->d is cycle-safe
    for _ in range(max_d_tries):

        d = int(rs.randint(n))
        if d in (a, b):
            continue
        if A[a, d]:  # multiedge
            continue

        # cycle safety for a->d
        if max_cycle_len == 3:
            if creates_triangle(a, d, neighbors_out, neighbors_in):
                continue
        else:
            if creates_short_cycle(a, d, neighbors_out, max_cycle_len):
                continue

        # need an existing edge c->d
        if not neighbors_in[d]:
            continue

        # try multiple c for this d
        for _ in range(max_c_tries):

            c = int(rs.choice(tuple(neighbors_in[d])))
            if c in (a, b, d):
                continue
            if A[c, b]:
                continue
            if len({a, b, c, d}) < 4:
                continue

            # cycle safety for c->b
            if max_cycle_len == 3:
                if creates_triangle(c, b, neighbors_out, neighbors_in):
                    continue
            else:
                if creates_short_cycle(c, b, neighbors_out, max_cycle_len):
                    continue

            # success
            return a, b, c, d, a, d, c, b

    return None

def run_constrained_ms_with_snapshots(
    A,
    n_accept_target=100000,
    n_attempt_max=200000,
    snapshots=10,
    snapshot_distrib='log',
    max_cycle_len=3,
    track_cycles=True,
    seed=None
):
    """
    Run constrained MS, collecting snapshots of A
    A: (n, n) binary adjacency

    Returns
    -------
    snaps : dict
        maps accepted swap counts to adjacency snapshots
    cycles : dict
        cycles counts at each snapshot
    info : dict
        run metadata
    """
    try:
        A = np.asarray(A)
    except TypeError as err:
        msg = ('A must be array_like. Received: {}.'.format(type(A)))
        raise TypeError(msg) from err
    
    rs = bct.get_rng(seed)

    W = A.copy().astype(float)
    A = (W != 0).astype(np.uint8).copy()
    np.fill_diagonal(W, 0.)
    np.fill_diagonal(A, 0)

    neighbors_out, neighbors_in = neighbors_from_dir_adj(A)

    if snapshot_distrib == 'linear':
        snapshot_accepts = np.linspace(10, n_accept_target, num=snapshots, dtype=int)
    elif snapshot_distrib == 'log':
        snapshot_accepts = np.logspace(1, np.log10(n_accept_target), num=snapshots, dtype=int)
    snapshot_accepts = set(snapshot_accepts)

    n_accept = 0
    n_attempt = 0
    snaps = {}

    if track_cycles:
        cycles = {'C3': [], 'C4': [], 'C5': [], 'swaps': []}
        C3, C4, C5 = global_cycle_counts_directed(A)
        cycles['C3'].append(C3)
        cycles['C4'].append(C4)
        cycles['C5'].append(C5)
        cycles['swaps'].append(0)

    step = max(1, n_accept_target // 100)
    pbar = tqdm(total=100, desc='constrained MS (accepted)')

    while n_accept < n_accept_target and n_attempt < n_attempt_max:

        n_attempt += 1

        prop = propose_targeted_ms_swap_directed(
            A, neighbors_out, neighbors_in,
            rs, max_cycle_len
        )
        if prop is None:
            continue
        a, b, c, d, u1, v1, u2, v2 = prop

        # Apply swap
        apply_swap_directed(
            A, W, neighbors_out, neighbors_in,
            a, b, c, d, u1, v1, u2, v2
        )

        n_accept += 1
        
        if n_accept % step == 0:
            pbar.update(1)

        if n_accept in snapshot_accepts:
            if track_cycles:
                C3, C4, C5 = global_cycle_counts_directed(A)
                # if any cycle count increases, stop
                if (C3 > cycles['C3'][-1] or
                    C4 > cycles['C4'][-1] or
                    C5 > cycles['C5'][-1]):
                    print(f"\nStopping early at {n_accept} accepts "
                          f"due to increase in cycle counts.")
                    break
                cycles['C3'].append(C3)
                cycles['C4'].append(C4)
                cycles['C5'].append(C5)
                cycles['swaps'].append(n_accept)
            snaps[n_accept] = W.copy()

    pbar.close()

    info = {
        'n_accept': n_accept,
        'n_attempt': n_attempt,
    }

    if track_cycles:
        return snaps, cycles, info
    else:
        return snaps, info