"""GF(2) rank / kernel of the SHA-256 sigma functions on 32-bit words.

Each of Sigma0, Sigma1, sigma0, sigma1 is GF(2)-linear (XOR of rotations / shifts), hence a
32x32 matrix over GF(2).  This script computes the rank, the kernel dimension and an explicit
kernel basis (every non-zero kernel vector is a counter-example to injectivity).  It is used
to decide which theorem to state in formal/Wp34/Sha256Words.lean before asking bv_decide.

Run:  .venv/bin/python formal/scripts/gf2_rank.py
"""
from __future__ import annotations
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
from wp34.sha256 import Sigma0, Sigma1, sigma0, sigma1, MASK  # noqa: E402

def columns(f):
    return [f(1 << i) & MASK for i in range(32)]  # image of basis vector e_i

def rank_and_kernel(f):
    cols = columns(f)
    # Gaussian elimination on the columns as 32-bit integers, tracking combinations for the kernel
    basis: list[tuple[int, int]] = []   # (pivot-reduced vector, combination mask over e_i)
    kernel: list[int] = []
    for i, v in enumerate(cols):
        comb = 1 << i
        for bv, bc in basis:
            if v & (1 << (bv.bit_length() - 1)):
                v ^= bv
                comb ^= bc
        if v == 0:
            kernel.append(comb)
        else:
            basis.append((v, comb))
            basis.sort(key=lambda t: -t[0])
    return len(basis), kernel

def main():
    for name, f in [("Sigma0", Sigma0), ("Sigma1", Sigma1), ("sigma0", sigma0), ("sigma1", sigma1)]:
        r, ker = rank_and_kernel(f)
        print(f"{name}: rank = {r}, kernel dimension = {32 - r}")
        for k in ker:
            assert f(k) & MASK == 0, (name, hex(k))
            print(f"  kernel vector x = {k:#010x}  ({name}(x) = 0, x != 0)")
        if not ker:
            # check that a random pair with equal image does not exist (bijective)
            print("  injective (kernel trivial) -> bijective on 32 bits")

if __name__ == "__main__" and "--lean" not in sys.argv:
    main()


# ---- Lean code generation (explicit inverses, pasted into formal/Wp34/Sha256Words.lean) ----
def inverse_columns(f):
    """Columns of the inverse matrix: inv(e_i) for i = 0..31 (Gauss-Jordan over GF(2))."""
    rows = [0] * 32
    for i in range(32):
        y = f(1 << i) & MASK
        for j in range(32):
            if (y >> j) & 1:
                rows[j] |= 1 << i
    I = [1 << j for j in range(32)]
    for bit in range(32):
        piv = next(i for i in range(bit, 32) if rows[i] >> bit & 1)
        rows[bit], rows[piv] = rows[piv], rows[bit]; I[bit], I[piv] = I[piv], I[bit]
        for i in range(32):
            if i != bit and rows[i] >> bit & 1:
                rows[i] ^= rows[bit]; I[i] ^= I[bit]
    assert rows == [1 << j for j in range(32)]
    def apply(x):
        return sum(((bin(I[j] & x).count("1") & 1) << j) for j in range(32))
    return [apply(1 << i) for i in range(32)], apply

def rotation_set(apply):
    """If inv is a XOR of rotations, return the rotation amounts (checked), else None."""
    from wp34.sha256 import rotr
    inv1 = apply(1)
    S = sorted(((32 - k) % 32) for k in range(32) if inv1 >> k & 1)
    import random
    for _ in range(2000):
        x = random.getrandbits(32); y = 0
        for j in S: y ^= rotr(x, j)
        if y != apply(x):
            return None
    return S

def lean_code():
    from wp34.sha256 import rotr
    out = []
    for name, f in [("Sigma0", Sigma0), ("Sigma1", Sigma1), ("sigma0", sigma0), ("sigma1", sigma1)]:
        cols, apply = inverse_columns(f)
        for _ in range(2000):
            x = __import__("random").getrandbits(32)
            assert apply(f(x) & MASK) == x and f(apply(x)) & MASK == x
        S = rotation_set(apply)
        if S is not None:
            body = " ^^^ ".join(f"rotr x {j}" for j in S)
            out.append(f"def {name}Inv (x : Word) : Word :=\n  {body}")
        else:
            body = " ^^^\n  ".join(f"bitMask x {i} &&& {c:#010x}#32" for i, c in enumerate(cols))
            out.append(f"def {name}Inv (x : Word) : Word :=\n  {body}")
    return "\n\n".join(out)

if __name__ == "__main__" and "--lean" in sys.argv:
    print(lean_code())