"""Reproduces the machine check of Section 4 of the note
"The partition exponent in the graded pieces of the Claude 3-SAT algorithm is
optimal for pieces of that shape".

Setting (Claude, "A 2^o(n) time algorithm for 3-SAT", hexagon:2610.00577, App. B).
V_S = Q[t_s : s in S], graded by weight; the basis is the monomials t^mu, mu a
partition into parts from S. nabla_{-k} f = t_k f and nabla_k f = k df/dt_k (k in S).
For a monomial a = t^mu and r in Z the operator a[r] of eq. (47) is computed by the
recursion (O2), with the modes l, k restricted to S:

    (t_n a')[r] = sum_{l in S, l >= n} C(l-1, n-1) nabla_{-l} a'[r-l]
                + sum_{k in S}        C(-k-1, n-1) a'[r+k] nabla_k,     1[r] = [r = 0] Id.

For S = all positive integers this is eq. (47) itself. Q_s(a, b) is eq. (48).

Tests, for all monomials a = t^mu (|mu| <= h), b = t^nu (|nu| = h), on V_E:
  Lemma 9 (a) Q_s(a,b) is unchanged when the factors are exchanged in every product;
          (b) Q_s(a,b) = 0 for |s| > E + h (checked for |s| up to E + h + 2);
          (c) sum_{|s| <= E+h} Q_s(a,b) = (a[-h] b) Id;
  Lemma 10 a^mu[-h] b^nu = h! [mu = nu] for partitions mu, nu of h into parts from S.

Expected: S = all positive integers passes everything (3878 triples (a, b, E) for
E, h <= 6); squares, cubes and odd numbers fail (a), (b), (c) already at a = b = t_1 on
V_0, e.g. for squares Q_{-3}(t_1, t_1) 1 = 4 in one order and 0 in the other, and
sum_{|s| <= 1} Q_s(t_1, t_1) 1 = -1 against the pairing t_1[-1] t_1 = 1.

Usage: python3 lemma9_check.py [max_weight]     (default 6; about two minutes)
"""
import sys
from collections import Counter
from functools import lru_cache
from math import comb, factorial


def gbinom(n, j):
    """C(n, j) for any integer n and j >= 0; 0 for j < 0."""
    if j < 0:
        return 0
    num = 1
    for i in range(j):
        num *= n - i
    return num // factorial(j)


def partitions(E, parts):
    """Partitions of E into the given parts, as nonincreasing tuples."""
    out = []

    def rec(rem, largest, acc):
        if rem == 0:
            out.append(tuple(acc))
            return
        for s in reversed(parts):
            if s <= largest and s <= rem:
                acc.append(s)
                rec(rem - s, s, acc)
                acc.pop()

    rec(E, E, [])
    return out


def add_into(out, vec, coef=1):
    for m, c in vec.items():
        out[m] = out.get(m, 0) + coef * c


def clean(vec):
    return {m: c for m, c in vec.items() if c}


class SparseFock:
    def __init__(self, modes):
        self.S = sorted(set(modes))
        self.Sset = set(self.S)
        self._basis = {}
        self.op_mono = lru_cache(maxsize=None)(self._op_mono)

    def basis(self, E):
        if E not in self._basis:
            self._basis[E] = partitions(E, self.S) if E >= 0 else []
        return self._basis[E]

    @staticmethod
    def times_t(k, vec):
        return {tuple(sorted(m + (k,), reverse=True)): c for m, c in vec.items()}

    @staticmethod
    def d_dt(k, mono):
        """nabla_k t^mono = k * mult_k * t^(mono - k)."""
        mult = mono.count(k)
        if not mult:
            return {}
        i = mono.index(k)
        return {mono[:i] + mono[i + 1:]: k * mult}

    def _op_mono(self, a, r, mono):
        """a[r] t^mono for a monomial a (tuple of parts), via (O2)."""
        if not a:
            return {mono: 1} if r == 0 else {}
        n, rest = a[0], a[1:]
        E = sum(mono)
        out = {}
        for l in self.S:                       # multiplications, applied last
            if l < n or l > E + r:
                continue
            c = gbinom(l - 1, n - 1)
            if c:
                add_into(out, self.times_t(l, self.op_mono(rest, r - l, mono)), c)
        for k in set(mono):                    # differentiations, applied first
            c = gbinom(-k - 1, n - 1)
            if c:
                for m1, c1 in self.d_dt(k, mono).items():
                    add_into(out, self.op_mono(rest, r + k, m1), c * c1)
        return clean(out)

    def op(self, avec, r, vec):
        out = {}
        for am, ac in avec.items():
            for m, c in vec.items():
                add_into(out, self.op_mono(am, r, m), ac * c)
        return clean(out)

    def Q(self, s, a, b, h, vec, exchanged=False):
        """Q_s(a, b) vec, eq. (48), for monomials a (weight <= h) and b (weight h)."""
        H = sum(a) + h
        out = {}
        for i in range(2 * h + 1):
            ai = comb(H, i)
            if not ai:
                continue
            if exchanged:
                t = self.op({b: 1}, h - s - i, self.op({a: 1}, s - h + i, vec))
            else:
                t = self.op({a: 1}, s - h + i, self.op({b: 1}, h - s - i, vec))
            add_into(out, t, (-1) ** i * ai)
        return clean(out)


def check_lemma9(F, max_w):
    counts = Counter()
    triples = 0
    for h in range(1, max_w + 1):
        As = [m for w in range(h + 1) for m in F.basis(w)]
        for E in range(max_w + 1):
            for a in As:
                for b in F.basis(h):
                    triples += 1
                    pairing = F.op({a: 1}, -h, {b: 1}).get((), 0)
                    for m in F.basis(E):
                        total = {}
                        for s in range(-(E + h) - 2, E + h + 3):
                            q1 = F.Q(s, a, b, h, {m: 1})
                            if q1 != F.Q(s, a, b, h, {m: 1}, exchanged=True):
                                counts["a"] += 1
                            if abs(s) > E + h and q1:
                                counts["b"] += 1
                            if abs(s) <= E + h:
                                add_into(total, q1)
                        if clean(total) != ({m: pairing} if pairing else {}):
                            counts["c"] += 1
    return triples, counts


def check_lemma10(F, h):
    def y(n):
        return {(j,): comb(n, j) for j in F.S if j <= n}

    def mul(u, v):
        out = {}
        for m1, c1 in u.items():
            for m2, c2 in v.items():
                add_into(out, {tuple(sorted(m1 + m2, reverse=True)): c1 * c2})
        return out

    def a_mu(mu):
        v = {(): 1}
        for p in mu:
            v = mul(v, {m: (-1) ** (p - 1) * c for m, c in y(p).items()})
        return v

    def z(mu):
        out = 1
        for k, m in Counter(mu).items():
            out *= k ** m * factorial(m)
        return out

    P = F.basis(h)
    wrong = sum(1 for mu in P for nu in P
                if F.op(a_mu(mu), -h, {nu: factorial(h) // z(nu)}).get((), 0) != (factorial(h) if mu == nu else 0))
    return wrong, len(P) ** 2


def main():
    max_w = int(sys.argv[1]) if len(sys.argv) > 1 else 6
    M = 4 * max_w + 8
    families = {
        "all positive integers (the paper)": list(range(1, M + 1)),
        "squares": [j * j for j in range(1, M) if j * j <= M],
        "cubes": [j ** 3 for j in range(1, M) if j ** 3 <= M],
        "odd numbers": [j for j in range(1, M + 1) if j % 2],
    }
    for name, modes in families.items():
        F = SparseFock(modes)
        triples, c = check_lemma9(F, max_w)
        print(f"S = {name}: Lemma 9 on {triples} triples (a, b, E), E, h <= {max_w}: "
              f"failures (a) {c['a']}, (b) {c['b']}, (c) {c['c']}")
        l10 = [check_lemma10(F, h) for h in range(1, max_w + 1)]
        print(f"   Lemma 10, h = 1..{max_w}: wrong entries {[w for w, _ in l10]} of {[n for _, n in l10]}")
        if name == "squares":
            t1 = (1,)
            one = {(): 1}
            print(f"   Q_(-3)(t_1, t_1) 1 = {F.Q(-3, t1, t1, 1, one).get((), 0)} in the order of (48), "
                  f"{F.Q(-3, t1, t1, 1, one, exchanged=True).get((), 0)} exchanged")
            tot = sum(F.Q(s, t1, t1, 1, one).get((), 0) for s in (-1, 0, 1))
            print(f"   sum_(|s|<=1) Q_s(t_1, t_1) 1 = {tot}; pairing t_1[-1] t_1 = {F.op({t1: 1}, -1, {t1: 1}).get((), 0)}")


if __name__ == "__main__":
    main()