"""Checks Lemma 6.2 (two counters) and Proposition 6.3 (Densest k-Subgraph by size and edges).

Lemma 6.2 computes the bivariate profile c[i][j] of a factor graph by evaluating the second counter at the
integer points u = 0, ..., D, working modulo 2^P with P = n + D + 1 and truncating the first counter at V^l.
Here the contractions are replaced by brute-force sums over {0,1}^n, reduced modulo 2^P exactly as the
contraction would be; everything after them is the recovery of the proof: forward differences, division by
t! with loss of v_2(t!) bits, and the change of basis from falling factorials to monomials.

Also checked: the recovered numbers equal the true counts; the claimed precision is enough; and it is needed
in the sense that the 2-adic valuation of t! really reaches t - s_2(t).
Run: python -I check_two_counters.py
"""
import itertools
import random

random.seed(20261030)


def v2(x):
    v = 0
    while x % 2 == 0:
        x //= 2
        v += 1
    return v


def falling_to_monomial(D):
    """coef[t][j]: coefficient of u^j in the falling factorial u(u-1)...(u-t+1), integers."""
    coef = [[0] * (D + 1) for _ in range(D + 1)]
    coef[0][0] = 1
    for t in range(1, D + 1):
        for j in range(D + 1):
            coef[t][j] += (coef[t - 1][j - 1] if j >= 1 else 0) - (t - 1) * coef[t - 1][j]
    return coef


def recover(values, D, n, P):
    """values[r] = g(r) mod 2^P for r = 0..D, g integer polynomial of degree <= D with coefficients in [0, 2^n].
    Returns the coefficients of g."""
    M = 1 << P
    binom = [[0] * (D + 2) for _ in range(D + 2)]
    for a in range(D + 2):
        binom[a][0] = 1
        for b in range(1, a + 1):
            binom[a][b] = binom[a - 1][b - 1] + binom[a - 1][b]
    fact = [1] * (D + 2)
    for t in range(1, D + 2):
        fact[t] = fact[t - 1] * t
    keep = n + 1  # bits kept at the end: modulus 2^(n+1) > 2^n
    b = []
    for t in range(D + 1):
        delta = sum((-1) ** (t - r) * binom[t][r] * values[r] for r in range(t + 1)) % M
        v = v2(fact[t])
        assert v == 0 or v <= t - 1, (t, v)
        assert delta % (1 << v) == 0, "forward difference must be divisible by 2^v2(t!)"
        odd = fact[t] >> v
        modulus = 1 << (P - v)
        assert P - v >= keep
        b.append(((delta >> v) * pow(odd, -1, modulus)) % modulus)
    # g = sum_t b_t u^{falling t}; monomial coefficients are integer combinations of the b_t
    coef = falling_to_monomial(D)
    out = [0] * (D + 1)
    mask = (1 << keep) - 1
    for t in range(D + 1):
        for j in range(D + 1):
            out[j] = (out[j] + coef[t][j] * b[t]) & mask
    return out


def eval_profile_graph(N, edges, u, P, ell):
    """Brute-force Z(Phi_u) over (Z/2^P Z)[V]/(V^ell) for the Densest k-Subgraph factor graph."""
    M = 1 << P
    z = [0] * ell
    powers = [pow(u, k, M) if not (u == 0 and k == 0) else 1 for k in range(N + 1)]
    for x in itertools.product((0, 1), repeat=N):
        e = sum(1 for a, b in edges if x[a] and x[b])
        k = sum(x)
        if e < ell:
            z[e] = (z[e] + powers[k]) % M
    return z


def check_dks(trials=150):
    for _ in range(trials):
        N = random.randint(1, 9)
        all_e = list(itertools.combinations(range(N), 2))
        edges = random.sample(all_e, random.randint(0, len(all_e)))
        Mm = len(edges)
        D = N
        P = N + D + 1
        ell = Mm + 1
        vals = [eval_profile_graph(N, edges, u, P, ell) for u in range(D + 1)]
        truth = [[0] * (D + 1) for _ in range(ell)]
        for x in itertools.product((0, 1), repeat=N):
            truth[sum(1 for a, b in edges if x[a] and x[b])][sum(x)] += 1
        for i in range(ell):
            rec = recover([vals[u][i] for u in range(D + 1)], D, N, P)
            assert rec == truth[i], (N, edges, i, rec, truth[i])
    print("Densest k-Subgraph profile recovered exactly with P = 2N+1 on", trials, "random graphs")


def check_general(trials=150):
    """Random two-counter factor graphs with hard constraints: tables with entries 0 or V^a U^b."""
    for _ in range(trials):
        n = random.randint(1, 7)
        nf = random.randint(1, 7)
        factors = []
        for _f in range(nf):
            r = random.randint(1, min(3, n))
            scope = tuple(random.sample(range(n), r))
            table = {}
            for beta in itertools.product((0, 1), repeat=r):
                if random.random() < 0.15:
                    table[beta] = None
                else:
                    table[beta] = (random.randint(0, 2), random.randint(0, 2))
            factors.append((scope, table))
        D = sum(max((t[1] for t in tab.values() if t), default=0) for _s, tab in factors)
        L = sum(max((t[0] for t in tab.values() if t), default=0) for _s, tab in factors) + 1
        P = n + D + 1
        Mod = 1 << P
        truth = [[0] * (D + 1) for _ in range(L)]
        for x in itertools.product((0, 1), repeat=n):
            a = b = 0
            ok = True
            for scope, tab in factors:
                t = tab[tuple(x[i] for i in scope)]
                if t is None:
                    ok = False
                    break
                a += t[0]
                b += t[1]
            if ok:
                truth[a][b] += 1
        vals = []
        for u in range(D + 1):
            z = [0] * L
            for x in itertools.product((0, 1), repeat=n):
                a = 0
                w = 1
                ok = True
                for scope, tab in factors:
                    t = tab[tuple(x[i] for i in scope)]
                    if t is None:
                        ok = False
                        break
                    a += t[0]
                    w = (w * (1 if t[1] == 0 else pow(u, t[1], Mod))) % Mod
                if ok:
                    z[a] = (z[a] + w) % Mod
            vals.append(z)
        for i in range(L):
            rec = recover([vals[u][i] for u in range(D + 1)], D, n, P)
            assert rec == truth[i], (n, factors, i, rec, truth[i])
    print("General two-counter factor graphs recovered exactly with P = n+D+1 on", trials, "instances")


def check_valuation():
    # v_2(t!) = t - (number of ones in t) <= t - 1 for t >= 1 (Legendre)
    f = 1
    for t in range(1, 2000):
        f *= t
        assert v2(f) == t - bin(t).count("1")
        assert v2(f) <= t - 1
    print("Legendre: v_2(t!) = t - s_2(t) <= t - 1 checked for t < 2000")


if __name__ == "__main__":
    check_valuation()
    check_dks()
    check_general()