"""Stage 5: stress tests for the column-deletion lemma P9 and the reduction
pipeline of Lemma 3.2 (results/derivations.md), written from scratch.

Test D1 (exhaustive, small X).  For 13 <= n <= 30 and every X with |X| <= 2
(|X| = 3 for n <= 20) supported on columns 11..n-1 (both layers), check the
strong form of P9: f(N[X]) = N[X'] and f(dX) = dX' as SETS (hence equal
sizes), where f deletes column 5 and shifts columns > 5 down by one.

Test D2 (negative control).  Same, but X is allowed to meet column 10 (gap
of only 10 empty columns).  Here the hypothesis of P9 fails and we count the
cases where f is undefined on dX (a neighbour of X sits in the deleted
column) or f(dX) != dX'.  A nonzero count shows the test has teeth.

Test D3 (pipeline on near-extremal sets).  For p in {16, 18} and several n0,
obtain an X with |dX| <= 11 from the project's own SAT formula with columns
0..9 forced empty (project encoder, model re-checked here by direct
counting), then (a) INSERT a column at position 5 repeatedly up to n0 + 25
and (b) DELETE it back down to n0; after every step the boundary is
recomputed from scratch in the new graph and must be unchanged, and the
final set must equal the original.
Output: results/robustness/deletion_lemma.json
"""
import itertools
import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.chdir(ROOT)
sys.path.insert(0, os.path.join(ROOT, ".pydeps"))
sys.path.insert(0, os.path.join(ROOT, "code"))


def adj(n, k=5):
    A = {}
    for i in range(n):
        A[("u", i)] = {("u", (i + 1) % n), ("u", (i - 1) % n), ("v", i)}
        A[("v", i)] = {("v", (i + k) % n), ("v", (i - k) % n), ("u", i)}
    return A


def closed_nbhd(A, X):
    return set(X) | {w for x in X for w in A[x]}


def boundary(A, X):
    return closed_nbhd(A, X) - set(X)


def f_del(w):
    layer, i = w
    if i == 5:
        return None
    return (layer, i if i < 5 else i - 1)


def f_ins(w):
    layer, i = w
    return (layer, i if i < 5 else i + 1)


def check_strong(n, X):
    A, A1 = adj(n), adj(n - 1)
    NX = closed_nbhd(A, X)
    if any(f_del(w) is None for w in NX):
        return False, "deleted vertex in N[X]"
    X1 = {f_del(w) for w in X}
    ok = {f_del(w) for w in NX} == closed_nbhd(A1, X1) and {f_del(w) for w in boundary(A, X)} == boundary(A1, X1)
    return ok, ""


def main():
    out = {}
    # D1
    d1_cases = d1_fail = 0
    for n in range(13, 31):
        verts = [(l, i) for i in range(11, n) for l in ("u", "v")]
        sizes = [1, 2] + ([3] if n <= 20 else [])
        for s in sizes:
            for X in itertools.combinations(verts, s):
                d1_cases += 1
                ok, _ = check_strong(n, set(X))
                if not ok:
                    d1_fail += 1
    out["D1"] = {"cases": d1_cases, "failures": d1_fail, "n_range": [13, 30], "sizes": "1,2 (and 3 for n<=20)"}
    # D2 negative control
    d2_cases = d2_undefined = d2_mismatch = d2_size_change = 0
    examples = []
    for n in range(13, 31):
        col10 = [("u", 10), ("v", 10)]
        others = [(l, i) for i in range(10, n) for l in ("u", "v")]
        for s in (1, 2):
            for X in itertools.combinations(others, s):
                if not any(w in col10 for w in X):
                    continue
                d2_cases += 1
                A, A1 = adj(n), adj(n - 1)
                NX = closed_nbhd(A, X)
                X1 = {f_del(w) for w in X}
                b1 = boundary(A1, X1)
                if any(f_del(w) is None for w in NX):
                    d2_undefined += 1
                    if len(b1) != len(boundary(A, X)):
                        d2_size_change += 1
                        if len(examples) < 3:
                            examples.append({"n": n, "X": sorted(X), "boundary_before": len(boundary(A, X)), "boundary_after": len(b1)})
                elif {f_del(w) for w in boundary(A, X)} != b1:
                    d2_mismatch += 1
    out["D2"] = {"cases": d2_cases, "f_undefined_on_N[X]": d2_undefined, "set_mismatch": d2_mismatch,
                 "boundary_size_changed": d2_size_change, "examples": examples}
    # D3 pipeline
    from sat_or_ilp_solver import boundary_clauses  # project encoder (only to obtain test instances)
    from pysat.solvers import Solver
    d3 = []
    for (p, n0) in [(16, 22), (16, 26), (16, 30), (18, 22), (18, 23), (18, 25), (20, 24)]:
        cl, x, z, nv = boundary_clauses(n0, 5, p, 11)
        # force columns 0..9 empty (rotation normalisation replaced by an explicit gap)
        cl = [c for c in cl if c != [x[0], x[n0]]]  # drop the project's x_{u0} v x_{v0} clause
        for i in range(10):
            cl.append([-x[i]])
            cl.append([-x[n0 + i]])
        with Solver(name="cadical195", bootstrap_with=cl) as sv:
            sat = sv.solve()
            model = sv.get_model() if sat else None
        rec = {"p": p, "n0": n0, "instance_found": bool(sat)}
        if sat:
            X = {("u" if w < n0 else "v", w if w < n0 else w - n0) for w in range(2 * n0) if model[x[w] - 1] > 0}
            A = adj(n0)
            b0 = len(boundary(A, X))
            rec["boundary_n0"] = b0
            rec["size_ok"] = len(X) == p
            rec["gap_ok"] = all(w[1] >= 10 for w in X)
            # insert up to n0+25
            cur, n = set(X), n0
            trail = []
            okall = True
            for step in range(25):
                cur = {f_ins(w) for w in cur}
                n += 1
                b = len(boundary(adj(n), cur))
                trail.append(b)
                okall &= (b == b0)
            rec["insert_boundaries_n0+1..n0+25_all_equal"] = okall
            okdel = True
            for step in range(25):
                cur = {f_del(w) for w in cur}
                n -= 1
                b = len(boundary(adj(n), cur))
                okdel &= (b == b0)
            rec["delete_back_all_equal"] = okdel
            rec["returned_to_original_set"] = (cur == X and n == n0)
            rec["ok"] = rec["size_ok"] and rec["gap_ok"] and okall and okdel and rec["returned_to_original_set"] and b0 <= 11
        d3.append(rec)
    out["D3"] = d3
    out["D1_ok"] = d1_fail == 0
    out["D3_ok"] = all(r["ok"] for r in d3 if r["instance_found"])
    out["D3_instances_found"] = sum(1 for r in d3 if r["instance_found"])
    json.dump(out, open("results/robustness/deletion_lemma.json", "w"), indent=1)
    print(json.dumps({k: out[k] for k in ("D1", "D2", "D1_ok", "D3_ok", "D3_instances_found")}, indent=None))
    for r in d3:
        print(r)


if __name__ == "__main__":
    main()
