"""Brute-force checks of the Bayes gadget (Lemma 4.1), the Bayes contexts (Lemma 4.2), the length gadget
(Lemma 4.3) and the Gibbs repetition numbers in the proof of Theorem 4.4.

Clauses are tuples of nonzero integers: +i is the literal z_i, -i its negation (variables are 1-based).
Run: python -I check_gadgets.py
"""
import itertools
import random
from fractions import Fraction
from math import comb

random.seed(20261030)


def satisfies(assign, clause):
    # assign: dict var -> 0/1
    return any((assign[abs(l)] == 1) == (l > 0) for l in clause)


def count_sat(clauses, nvars):
    total = 0
    sols = []
    for bits in itertools.product((0, 1), repeat=nvars):
        a = {i + 1: bits[i] for i in range(nvars)}
        if all(satisfies(a, c) for c in clauses):
            total += 1
            sols.append(a)
    return total, sols


def random_kcnf(n, m, k):
    clauses = set()
    tries = 0
    while len(clauses) < m and tries < 1000:
        tries += 1
        r = random.randint(1, min(k, n))
        vs = random.sample(range(1, n + 1), r)
        clauses.add(tuple(sorted(v if random.random() < 0.5 else -v for v in vs)))
    return list(clauses)


def bayes_gadget(F, n):
    """Gamma_F on variables 1..n (z), n+1 (a), n+2 (u)."""
    a, u = n + 1, n + 2
    G = [tuple(sorted(c + (-a,))) for c in F]
    G += [(a, -i) for i in range(1, n + 1)]
    G += [(a, -u)]
    return G, a


# ---------------- literal disjunctions ----------------
def ld_label(sigma, x):
    # sigma, x: dict coord -> value; x has entries +-1 on its support
    return int(any(sigma[i] == v for i, v in x.items()))


def clause_to_ld(clause):
    return {abs(l): (1 if l > 0 else -1) for l in clause}


def posterior_predictive_ld(context, query, dim):
    vs = []
    for sig in itertools.product((-1, 1), repeat=dim):
        sigma = {i + 1: sig[i] for i in range(dim)}
        if all(ld_label(sigma, x) == y for x, y in context):
            vs.append(sigma)
    if not vs:
        return None, 0
    pos = sum(ld_label(s, query) for s in vs)
    return Fraction(pos, len(vs)), len(vs)


# ---------------- binarized perceptrons via the gadget of [Li26, Lemma 3.3] ----------------
def perceptron_gadget(G, nvars, k):
    p = max(1, k - 2)
    ex = []
    for c in G:
        r = len(c)
        x = {abs(l): (1 if l > 0 else -1) for l in c}
        if r >= 2:
            for j in range(1, r - 1):
                x[nvars + j] = 1
        else:
            x[nvars + 1] = -1
        ex.append((x, 1))
    for j in range(1, p + 1):
        ex.append(({nvars + j: 1}, 1))
    return ex, nvars + p


def bp_label(sigma, x):
    return int(sum(sigma[i] * v for i, v in x.items()) >= 0)


def posterior_predictive_bp(context, query, dim):
    vs = []
    for sig in itertools.product((-1, 1), repeat=dim):
        sigma = {i + 1: sig[i] for i in range(dim)}
        if all(bp_label(sigma, x) == y for x, y in context):
            vs.append(sigma)
    if not vs:
        return None, 0
    pos = sum(bp_label(s, query) for s in vs)
    return Fraction(pos, len(vs)), len(vs)


def check_bayes_gadget(trials=300):
    for _ in range(trials):
        n = random.randint(1, 5)
        k = random.randint(2, 3)
        m = random.randint(0, 7)
        F = random_kcnf(n, m, k)
        nF, _ = count_sat(F, n)
        G, a = bayes_gadget(F, n)
        nG, solsG = count_sat(G, n + 2)
        assert nG == 2 * nF + 1, (F, nF, nG)
        assert sum(s[a] for s in solsG) == 2 * nF
        assert all(len(c) <= k + 1 for c in G)
        # LD context
        ctx = [(clause_to_ld(c), 1) for c in G]
        assert all(len(x) <= k + 1 for x, _ in ctx)
        p, size = posterior_predictive_ld(ctx, {a: 1}, n + 2)
        assert size == 2 * nF + 1 and p == Fraction(2 * nF, 2 * nF + 1), (F, p)
        # BP context: Gamma_F is a (k+1)-CNF, padding p = k-1
        ctxb, dim = perceptron_gadget(G, n + 2, k + 1)
        assert dim == n + k + 1
        assert len(ctxb) == len(F) + n + k
        assert all(len(x) <= 2 * k for x, _ in ctxb)
        pb, sizeb = posterior_predictive_bp(ctxb, {a: 1}, dim)
        assert sizeb == 2 * nF + 1 and pb == Fraction(2 * nF, 2 * nF + 1), (F, pb)
        if p > 0:
            assert p >= Fraction(2, 3)
            assert Fraction(nF) == p / (2 * (1 - p))
    print("Bayes gadget and Bayes contexts (LD and BP): OK on", trials, "random formulas")


# ---------------- length gadget ----------------
def split_to_3cnf(F, n):
    """Equisatisfiable 3-CNF by chains; returns (clauses, number of variables)."""
    nv = n
    out = []
    for c in F:
        c = list(c)
        r = len(c)
        if r <= 3:
            out.append(tuple(c))
            continue
        ts = list(range(nv + 1, nv + r - 2))
        nv += r - 3
        out.append((c[0], c[1], ts[0]))
        for j in range(1, r - 3):
            out.append((-ts[j - 1], c[j + 1], ts[j]))
        out.append((-ts[-1], c[-2], c[-1]))
    return out, nv


def length_gadget(F3, nv):
    """Gamma' of Lemma 4.3 on z (1..nv), y_C, a, u."""
    cur = nv
    ys = []
    for c in F3:
        if len(c) <= 2:
            ys.append((c, None))
        else:
            cur += 1
            ys.append((c, cur))
    a, u = cur + 1, cur + 2
    G = []
    for c, y in ys:
        if y is None:
            G.append(tuple(c) + (-a,))
        else:
            l1, l2, l3 = c
            G.append((-a, l1, y))
            G.append((-y, l2, l3))
            G.append((a, -y))
    for i in range(1, nv + 1):
        G.append((a, -i))
    G.append((a, -u))
    return G, cur + 2, a


def check_length_gadget(trials=200):
    for _ in range(trials):
        n = random.randint(1, 5)
        m = random.randint(1, 5)
        F = []
        for _ in range(m):
            r = random.randint(1, min(5, n))
            vs = random.sample(range(1, n + 1), r)
            F.append(tuple(v if random.random() < 0.5 else -v for v in vs))
        L = sum(len(c) for c in F)
        nF, _ = count_sat(F, n)
        # delete variables that do not occur, as in step <1>1 of the proof (satisfiability is unchanged)
        occ = sorted({abs(l) for c in F for l in c})
        ren = {v: i + 1 for i, v in enumerate(occ)}
        Fr = [tuple((1 if l > 0 else -1) * ren[abs(l)] for l in c) for c in F]
        F3, nv = split_to_3cnf(Fr, len(occ))
        n3, _ = count_sat(F3, nv)
        assert (n3 > 0) == (nF > 0)
        G, N, a = length_gadget(F3, nv)
        assert all(len(c) <= 3 for c in G)
        assert N <= 3 * L + 2 and len(G) <= 5 * L + 1, (N, L, len(G))
        nG, sols = count_sat(G, N)
        assert sum(1 for s in sols if s[a] == 0) == 1
        A = sum(1 for s in sols if s[a] == 1)
        assert (A == 0) == (nF == 0) and (A == 0 or A >= 2)
        if N <= 12:
            ctx = [(clause_to_ld(c), 1) for c in G]
            p, _ = posterior_predictive_ld(ctx, {a: 1}, N)
            assert p == Fraction(A, A + 1)
            assert (p == 0) if nF == 0 else (p >= Fraction(2, 3))
    print("Length gadget (Lemma 4.3): OK on", trials, "random formulas")


def check_gibbs_numbers():
    def tail_ge(nn, q, t):
        return sum(comb(nn, i) * q ** i * (1 - q) ** (nn - i) for i in range(t, nn + 1))

    p1 = tail_ge(12, Fraction(1, 6), 4)
    p2 = 1 - tail_ge(12, Fraction(1, 2), 4)
    assert p1 < Fraction(1, 7), float(p1)
    assert p2 == Fraction(299, 4096) and p2 < Fraction(1, 13)
    print("Gibbs repetition numbers: Pr[Bin(12,1/6)>=4] = %.5f < 1/7, Pr[Bin(12,1/2)<=3] = 299/4096" % float(p1))


if __name__ == "__main__":
    check_bayes_gadget()
    check_length_gadget()
    check_gibbs_numbers()