"""Zero-forcing closure oracle (definition D3) with complete chronology, replay,
and the order-independence test T2 (empirical companion of Proposition P1).

A coloring is a set of black vertices. A black vertex b with exactly one white
neighbour w may force w. closure() performs forces until none is available and
returns (final_black_set, chronology); the chronology is the list of forces
(b, w) in the order performed. The scan order used to pick the next force is a
parameter (fixed increasing order, or a seeded random permutation at every
step); by P1 the final set does not depend on it.
"""
import json
import random
import sys

import common  # noqa: F401
from graph_constructor import adjacency, label, labels_to_ids


def closure(adj, S, order="fixed", rng=None):
    """adj: adjacency lists; S: iterable of vertex ids.
    order: 'fixed' (lowest-index available force first) or 'random' (a fresh
    random permutation of the candidate black vertices at every step, using rng)."""
    N = len(adj)
    black = [False] * N
    for w in S:
        black[w] = True
    chron = []
    while True:
        candidates = range(N) if order == "fixed" else rng.sample(range(N), N)
        forced = False
        for b in candidates:
            if not black[b]:
                continue
            white = [x for x in adj[b] if not black[x]]
            if len(white) == 1:
                w = white[0]
                black[w] = True
                chron.append((b, w))
                forced = True
                break  # restart the scan: 'fixed' order is then deterministic
        if not forced:
            break
    final = {w for w in range(N) if black[w]}
    return final, chron


def closure_fast(adjmask, Smask):
    """Bitmask closure (same rule); used by solvers for speed. Returns final mask."""
    B = Smask
    N = len(adjmask)
    changed = True
    while changed:
        changed = False
        rest = B
        while rest:
            low = rest & -rest
            b = low.bit_length() - 1
            rest ^= low
            white = adjmask[b] & ~B
            if white and (white & (white - 1)) == 0:
                B |= white
                changed = True
    return B


def is_zero_forcing(adj, S):
    final, _ = closure(adj, S)
    return len(final) == len(adj)


def replay(adj, S, chron):
    """Replays a chronology, checking legality of every force. Returns
    (ok, final_black_set, first_bad_index)."""
    N = len(adj)
    black = [False] * N
    for w in S:
        black[w] = True
    for idx, (b, w) in enumerate(chron):
        if not black[b] or black[w]:
            return False, {x for x in range(N) if black[x]}, idx
        white = [x for x in adj[b] if not black[x]]
        if white != [w]:
            return False, {x for x in range(N) if black[x]}, idx
        black[w] = True
    return True, {x for x in range(N) if black[x]}, None


def order_independence_test(n, k, seeds_per_n=200, orders_per_seed=20, sizes=(3, 8), seed=20261030):
    """T2: for random S, compare the final black set over random force orders."""
    adj = adjacency(n, k)
    N = 2 * n
    rng = random.Random(seed * 1000 + n)
    passed = 0
    total = 0
    for _ in range(seeds_per_n):
        s = rng.randint(sizes[0], sizes[1])
        S = rng.sample(range(N), s)
        ref, _ = closure(adj, S)
        finals = set()
        for _ in range(orders_per_seed):
            fin, _ = closure(adj, S, order="random", rng=rng)
            finals.add(frozenset(fin))
        total += 1
        if len(finals) == 1 and frozenset(ref) in finals:
            passed += 1
    return passed, total


if __name__ == "__main__":
    k = 5
    report = []
    allok = True
    for n in range(11, 17):
        p, t = order_independence_test(n, k)
        print("T2 n=%d: %d/%d seed sets had an order-independent closure" % (n, p, t))
        report.append({"n": n, "k": k, "passed": p, "total": t})
        allok &= p == t
    print("T2 ALL OK" if allok else "T2 FAILURE")
    with open("results/t2_order_independence.json", "w") as f:
        json.dump(report, f, indent=1)