"""Stage 5: heuristic attacks on Theorem 5.1 and on the boundary certificates.

Attack A (theorem): simulated annealing over 11-subsets S of V(P(n,5)),
maximising |cl(S)|, for 23 <= n <= 60 and a sample of larger n.  Any S with
cl(S) = V would be a counterexample to Z(P(n,5)) = 12.  Positive controls:
(n, s) = (22, 11), (21, 10), (13, 9), (11, 6), where forcing sets exist.

Attack B (certificates): simulated annealing over 18-subsets X, minimising
|dX|, for 26 <= n <= 60 and a sample up to n = 200.  Any X with |dX| <= 11
would contradict the DRAT-certified unsatisfiability of F(n,18) (n <= 154)
or Lemma 3.2 (n > 154).  Positive controls: n = 22, 23, 25 (beta_18 = 10,
11, 10 in results/tables/boundary_table.csv) and p = 20 at n = 24 (11).

Everything here is written from scratch (bitmask closure, no project
imports).  Seeds are fixed and recorded.  Output:
results/robustness/local_search_attacks.json
"""
import json
import math
import os
import random
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.chdir(ROOT)


def masks(n, k):
    N = 2 * n
    nb = [0] * N
    for i in range(n):
        u, v = i, n + i
        for t in ((i + 1) % n, (i - 1) % n, v):
            nb[u] |= 1 << t
        for t in (n + (i + k) % n, n + (i - k) % n, u):
            nb[v] |= 1 << t
    return N, nb


def closure_size(N, nb, S):
    B = S
    full = (1 << N) - 1
    changed = True
    while changed and B != full:
        changed = False
        rest = B
        while rest:
            low = rest & -rest
            b = low.bit_length() - 1
            rest ^= low
            white = nb[b] & ~B
            if white and white & (white - 1) == 0:
                B |= white
                changed = True
    return bin(B).count("1")


def boundary_size(nb, X):
    NX = 0
    rest = X
    while rest:
        low = rest & -rest
        NX |= nb[low.bit_length() - 1]
        rest ^= low
    return bin(NX & ~X).count("1")


def anneal(N, score, s, iters, restarts, rng, target, maximise=True):
    """Generic swap-move annealing over s-subsets of range(N). Returns best score and set."""
    best_val, best_set = None, None
    sign = 1 if maximise else -1
    for _ in range(restarts):
        cur = rng.sample(range(N), s)
        curmask = sum(1 << w for w in cur)
        cur_val = score(curmask)
        if best_val is None or sign * cur_val > sign * best_val:
            best_val, best_set = cur_val, curmask
        T0, T1 = 1.5, 0.05
        for it in range(iters):
            T = T0 * (T1 / T0) ** (it / max(1, iters - 1))
            i = rng.randrange(s)
            w_out = cur[i]
            while True:
                w_in = rng.randrange(N)
                if not (curmask >> w_in) & 1:
                    break
            newmask = (curmask & ~(1 << w_out)) | (1 << w_in)
            val = score(newmask)
            d = sign * (val - cur_val)
            if d >= 0 or rng.random() < math.exp(d / T):
                cur[i] = w_in
                curmask, cur_val = newmask, val
                if sign * val > sign * best_val:
                    best_val, best_set = val, newmask
                    if (maximise and val >= target) or (not maximise and val <= target):
                        return best_val, best_set
    return best_val, best_set


def mask_to_labels(n, m):
    return ["%s%d" % ("u" if w < n else "v", w if w < n else w - n) for w in range(2 * n) if (m >> w) & 1]


def main():
    seed = 20261030
    rng = random.Random(seed)
    out = {"seed": seed, "attack_A": [], "attack_A_controls": [], "attack_B": [], "attack_B_controls": []}
    t0 = time.time()
    # ---- Attack A controls
    for (n, s) in [(11, 6), (13, 9), (21, 10), (22, 11)]:
        N, nb = masks(n, 5)
        val, S = anneal(N, lambda m: closure_size(N, nb, m), s, 20000, 6, rng, N, True)
        out["attack_A_controls"].append({"n": n, "s": s, "best_closure": val, "N": N, "forcing_found": val == N,
                                         "S": mask_to_labels(n, S) if val == N else None})
    # ---- Attack A: 11-sets for n = 23..60 and sample
    ns = list(range(23, 61)) + [70, 80, 96, 120, 154, 155, 200]
    for n in ns:
        N, nb = masks(n, 5)
        val, S = anneal(N, lambda m: closure_size(N, nb, m), 11, 15000, 4, rng, N, True)
        rec = {"n": n, "s": 11, "best_closure": val, "N": N, "forcing_found": val == N}
        if val == N:
            rec["COUNTEREXAMPLE_S"] = mask_to_labels(n, S)
        out["attack_A"].append(rec)
    tA = time.time() - t0
    # ---- Attack B controls
    for (n, p, expect) in [(22, 18, 10), (23, 18, 11), (25, 18, 10), (24, 20, 11), (24, 18, 12)]:
        N, nb = masks(n, 5)
        val, X = anneal(N, lambda m: boundary_size(nb, m), p, 40000, 6, rng, expect, False)
        out["attack_B_controls"].append({"n": n, "p": p, "expected_beta_from_table": expect, "best_boundary": val,
                                         "X": mask_to_labels(n, X)})
    # ---- Attack B: 18-sets, n = 26..60 and sample
    ns = list(range(26, 61)) + [70, 80, 96, 120, 154, 155, 160, 200]
    for n in ns:
        N, nb = masks(n, 5)
        val, X = anneal(N, lambda m: boundary_size(nb, m), 18, 40000, 5, rng, 11, False)
        rec = {"n": n, "p": 18, "best_boundary": val, "violation": val <= 11}
        if val <= 11:
            rec["COUNTEREXAMPLE_X"] = mask_to_labels(n, X)
        out["attack_B"].append(rec)
    out["wall_seconds_A"] = round(tA, 1)
    out["wall_seconds_total"] = round(time.time() - t0, 1)
    out["attack_A_any_counterexample"] = any(r["forcing_found"] for r in out["attack_A"])
    out["attack_B_any_violation"] = any(r["violation"] for r in out["attack_B"])
    out["attack_A_controls_all_found"] = all(r["forcing_found"] for r in out["attack_A_controls"])
    out["attack_B_controls_reached_table_value"] = all(r["best_boundary"] == r["expected_beta_from_table"] for r in out["attack_B_controls"])
    json.dump(out, open("results/robustness/local_search_attacks.json", "w"), indent=1)
    print("A controls found:", out["attack_A_controls_all_found"], [(r["n"], r["best_closure"], r["N"]) for r in out["attack_A_controls"]])
    print("A counterexample:", out["attack_A_any_counterexample"], "best closures:", [(r["n"], r["best_closure"], r["N"]) for r in out["attack_A"]])
    print("B controls:", [(r["n"], r["p"], r["best_boundary"], r["expected_beta_from_table"]) for r in out["attack_B_controls"]])
    print("B violation:", out["attack_B_any_violation"], "best boundaries:", [(r["n"], r["best_boundary"]) for r in out["attack_B"]])
    print("wall", out["wall_seconds_total"])


if __name__ == "__main__":
    main()
