# Brute-force checks for GL_{2m}/Sp_{2m}: fibres of the convolution map over x_nu,
# and orbit counts mod t^3 for GL_4.  All computations over F_p, p prime.
import itertools, sys

def subspaces(n, k, p):
    """All k-dim subspaces of F_p^n as tuples of basis vectors in reduced row echelon form."""
    # enumerate RREF matrices
    res = []
    for pivots in itertools.combinations(range(n), k):
        free_pos = []
        for r, pc in enumerate(pivots):
            for c in range(pc + 1, n):
                if c not in pivots:
                    free_pos.append((r, c))
        for vals in itertools.product(range(p), repeat=len(free_pos)):
            M = [[0] * n for _ in range(k)]
            for r, pc in enumerate(pivots):
                M[r][pc] = 1
            for (r, c), v in zip(free_pos, vals):
                M[r][c] = v
            res.append(tuple(tuple(row) for row in M))
    return res

def alt_form(blocks, p):
    """Form on O^{2m} with omega(e_{2i-1}, e_{2i}) = t^{a_i}; returns list of coefficient matrices omega_r (r = 0..maxdeg)."""
    m = len(blocks); n = 2 * m
    maxd = max(blocks)
    mats = [[[0] * n for _ in range(n)] for _ in range(maxd + 1)]
    for i, a in enumerate(blocks):
        mats[a][2 * i][2 * i + 1] = 1
        mats[a][2 * i + 1][2 * i] = -1 % p
    return mats

def bil(M, u, v, p):
    return sum(u[i] * M[i][j] * v[j] for i in range(len(u)) for j in range(len(v))) % p

def fibre_count(blocks, k, p):
    """# of W in Gr(k, t^{-1}L0/L0) with omega perfect on L0 + W.
    Perfect <=> integral (index q^k and v(Pf)=k=sum blocks) <=> omega_0(W, F^n)=0 and omega_1|_W = 0."""
    assert sum(blocks) == k
    mats = alt_form(blocks, p); n = 2 * len(blocks)
    cnt = 0
    basis_all = [tuple(1 if j == i else 0 for j in range(n)) for i in range(n)]
    for W in subspaces(n, k, p):
        ok = True
        for w in W:
            if any(bil(mats[0], w, e, p) for e in basis_all):
                ok = False; break
        if ok and len(mats) > 1:
            for w1 in W:
                for w2 in W:
                    if bil(mats[1], w1, w2, p):
                        ok = False; break
                if not ok: break
        if ok: cnt += 1
    return cnt

def lagr(n1, q):
    r = 1
    for i in range(1, n1 + 1): r *= (1 + q ** i)
    return r

print("== GL_4: fibre of Gr(2,4)-convolution over x_nu ==")
for p in (2, 3):
    for blocks in ((1, 1), (2, 0)):
        n1 = blocks.count(1)
        print(f"p={p} nu={blocks}: count={fibre_count(blocks, 2, p)}  predicted prod(1+q^i,i<=n1)={lagr(n1, p)}")

print("== GL_6: fibre of Gr(k,6)-convolution over x_nu, p=2 ==")
p = 2
for k in (2, 3, 4):
    for blocks in itertools.product((0, 1, 2), repeat=3):
        if sum(blocks) == k and list(blocks) == sorted(blocks, reverse=True):
            n1 = blocks.count(1)
            print(f"k={k} nu={blocks}: count={fibre_count(blocks, k, p)}  predicted={lagr(n1, p)}")

print("== GL_4: T_(2,0,0,0)-fibre: lattices L0 + O t^{-2} v, v primitive in (O/t^2)^4 ==")
def cyclic_fibre(blocks, p):
    mats = alt_form(blocks, p); n = 4
    # v = v0 + t v1, v0 != 0 ; L' = L0 + O t^{-2} v. integrality: omega(v, L0) in t^2 O
    # omega(v,e) = sum_r t^r omega_r(v0,e) + t^{r+1} omega_r(v1,e); need coefficients of t^0, t^1 vanish
    seen = set()
    for v0 in itertools.product(range(p), repeat=n):
        if not any(v0): continue
        for v1 in itertools.product(range(p), repeat=n):
            ok = True
            for e in [tuple(1 if j == i else 0 for j in range(n)) for i in range(n)]:
                c0 = bil(mats[0], v0, e, p)
                c1 = (bil(mats[1], v0, e, p) + bil(mats[0], v1, e, p)) % p if len(mats) > 1 else bil(mats[0], v1, e, p)
                if c0 or c1: ok = False; break
            if ok:
                # normalise v modulo units of O/t^2: canonical rep = the O/t^2-line through v
                line = set()
                for u0 in range(1, p):
                    for u1 in range(p):
                        w0 = tuple((u0 * a) % p for a in v0)
                        w1 = tuple((u0 * b + u1 * a) % p for a, b in zip(v0, v1))
                        line.add((w0, w1))
                seen.add(frozenset(line))
    return len(seen)
for p in (2, 3):
    print(f"p={p}: nu=(2,0): {cyclic_fibre((2,0),p)} (predicted q(q+1)={p*(p+1)}); nu=(1,1): {cyclic_fibre((1,1),p)} (predicted 0)")

print("== GL_4: orbit counts mod t^3 over F_2 by type (v(Pf), min valuation) ==")
p = 2; N = 3
def pf(w):  # w = dict of 6 entries as polynomials (lists length N) ; Pf = w12 w34 - w13 w24 + w14 w23
    def mul(a, b):
        c = [0] * N
        for i in range(N):
            for j in range(N - i):
                c[i + j] = (c[i + j] + a[i] * b[j]) % p
        return c
    def add(a, b, s=1): return [(x + s * y) % p for x, y in zip(a, b)]
    return add(add(mul(w[0], w[5]), mul(w[1], w[4]), -1), mul(w[2], w[3]))
from collections import Counter
cnt = Counter()
for coeffs in itertools.product(range(p), repeat=6 * N):
    w = [list(coeffs[6 * r:6 * r + 6]) for r in range(N)]  # w[r] = coefficient of t^r (6 entries)
    entries = [[w[r][e] for r in range(N)] for e in range(6)]
    P = pf(entries)
    vP = next((i for i, c in enumerate(P) if c), N)
    minv = min((next((i for i, c in enumerate(ent) if c), N) for ent in entries))
    cnt[(vP, minv)] += 1
for key in sorted(cnt): print(key, cnt[key])
q = 2
print("predicted #(2,2,0,0) = q^10 (q-1)(q^3-1)(q^2+1) =", q**10*(q-1)*(q**3-1)*(q**2+1), "; #(1,1,1,1) = q^8 (q-1)(q^3-1) =", q**8*(q-1)*(q**3-1))