# S1-4 checks for GL_{2m}/Sp_{2m}:
# (a) dim_k {X in gl_2m(k[t]/t^N) : X D_a + D_a X^T = 0} = (2m^2+m)N + sum_j (4j-3) a_j,
#     by linear algebra over F_p (p = 3, 5, 7), all labels a >= 0 with a_1 < N, m = 1..3.
# (b) Cayley transform: for random X in that space with 1-X invertible,
#     g = (1-X)^{-1}(1+X) satisfies g D_a g^T = D_a mod t^N.
# (c) def(nu) = n_1(nu) for nu in {0,1,2}^m, |nu| = k, m <= 8, using only the orbit-dimension formula.
import itertools, random, sys
random.seed(1)

def labels(m, N):
    for a in itertools.product(range(N), repeat=m):
        if all(a[i] >= a[i+1] for i in range(m-1)):
            yield a

# matrices over R = F_p[t]/t^N : entries are lists of length N
def Dmat(a, N, p):
    n = 2*len(a)
    D = [[[0]*N for _ in range(n)] for _ in range(n)]
    for j, aj in enumerate(a):
        if aj < N:
            D[2*j][2*j+1][aj] = 1
            D[2*j+1][2*j][aj] = p-1
    return D

def rmul(x, y, N, p):
    r = [0]*N
    for i, c in enumerate(x):
        if c:
            for j in range(N-i):
                if y[j]:
                    r[i+j] = (r[i+j] + c*y[j]) % p
    return r

def mmul(A, B, N, p):
    n = len(A)
    C = [[[0]*N for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for k in range(n):
            if any(A[i][k]):
                for j in range(n):
                    if any(B[k][j]):
                        pr = rmul(A[i][k], B[k][j], N, p)
                        C[i][j] = [(u+v) % p for u, v in zip(C[i][j], pr)]
    return C

def transpose(A):
    n = len(A)
    return [[A[j][i] for j in range(n)] for i in range(n)]

def rank_mod(rows, p):
    rows = [r[:] for r in rows]
    rk = 0; ncol = len(rows[0]) if rows else 0
    for c in range(ncol):
        piv = None
        for i in range(rk, len(rows)):
            if rows[i][c] % p:
                piv = i; break
        if piv is None: continue
        rows[rk], rows[piv] = rows[piv], rows[rk]
        inv = pow(rows[rk][c], p-2, p)
        rows[rk] = [(v*inv) % p for v in rows[rk]]
        for i in range(len(rows)):
            if i != rk and rows[i][c] % p:
                f = rows[i][c]
                rows[i] = [(u - f*v) % p for u, v in zip(rows[i], rows[rk])]
        rk += 1
    return rk, rows

def lie_dim(a, N, p):
    n = 2*len(a); D = Dmat(a, N, p)
    nv = n*n*N
    # linear map X -> X D + D X^T ; columns indexed by (i,j,e) basis matrices t^e E_ij
    cols = []
    for i in range(n):
        for j in range(n):
            for e in range(N):
                X = [[[0]*N for _ in range(n)] for _ in range(n)]
                X[i][j][e] = 1
                Y = mmul(X, D, N, p); Z = mmul(D, transpose(X), N, p)
                vec = []
                for r in range(n):
                    for s in range(n):
                        vec += [(u+v) % p for u, v in zip(Y[r][s], Z[r][s])]
                cols.append(vec)
    # rank of the map = rank of matrix with these columns
    rk, _ = rank_mod(cols, p)
    return nv - rk, cols

def formula(a, N):
    m = len(a)
    return (2*m*m+m)*N + sum((4*(j+1)-3)*aj for j, aj in enumerate(a))

ok = True
for p in [3, 5, 7]:
    for m in [1, 2, 3]:
        for N in ([2, 3, 4] if m < 3 else [2, 3]):
            if p == 7 and m == 3 and N == 3: continue
            for a in labels(m, N):
                dL, _ = lie_dim(a, N, p)
                if dL != formula(a, N):
                    ok = False; print('MISMATCH', p, m, N, a, dL, formula(a, N))
    print('p=%d: Lie-algebra dimension formula checked, all ok so far: %s' % (p, ok)); sys.stdout.flush()

# (b) Cayley transform
def kernel_basis(a, N, p):
    n = 2*len(a); D = Dmat(a, N, p)
    _, cols = lie_dim(a, N, p)
    # solve sum z_k cols[k] = 0 : build matrix with rows = equations
    neq = len(cols[0]); nv = len(cols)
    A = [[cols[k][r] for k in range(nv)] for r in range(neq)]
    rk, R = rank_mod(A, p)
    piv = []
    for r in range(rk):
        c = next(c for c in range(nv) if R[r][c] % p)
        piv.append(c)
    free = [c for c in range(nv) if c not in piv]
    basis = []
    for f in free:
        z = [0]*nv; z[f] = 1
        for r, c in enumerate(piv):
            z[c] = (-R[r][f]) % p
        basis.append(z)
    return basis

def vec_to_mat(z, n, N):
    X = [[[0]*N for _ in range(n)] for _ in range(n)]
    k = 0
    for i in range(n):
        for j in range(n):
            for e in range(N):
                X[i][j][e] = z[k]; k += 1
    return X

def inv_mat(A, N, p):
    # Gauss-Jordan over R = F_p[t]/t^N for A with det unit
    n = len(A)
    M = [[A[i][j][:] for j in range(n)] + [([1]+[0]*(N-1)) if i == j else [0]*N for j in range(n)] for i in range(n)]
    def runit_inv(x):
        g = [0]*N; g[0] = pow(x[0], p-2, p)
        for k in range(1, N):
            s = sum(x[i]*g[k-i] for i in range(1, k+1))
            g[k] = (-s*g[0]) % p
        return g
    for c in range(n):
        piv = next(i for i in range(c, n) if M[i][c][0] % p)
        M[c], M[piv] = M[piv], M[c]
        u = runit_inv(M[c][c])
        M[c] = [rmul(u, x, N, p) for x in M[c]]
        for i in range(n):
            if i != c and any(M[i][c]):
                f = M[i][c]
                M[i] = [[(x - y) % p for x, y in zip(M[i][j], rmul(f, M[c][j], N, p))] for j in range(2*n)]
    return [row[n:] for row in M]

cay_ok = True; tested = 0
for p in [3, 5]:
    for (m, N, a) in [(2, 3, (2, 0)), (2, 3, (1, 1)), (3, 3, (2, 1, 0)), (3, 4, (3, 1, 1))]:
        n = 2*m; D = Dmat(a, N, p); B = kernel_basis(a, N, p)
        for trial in range(5):
            z = [0]*len(B[0])
            for b in B:
                c = random.randrange(p)
                z = [(u + c*v) % p for u, v in zip(z, b)]
            X = vec_to_mat(z, n, N)
            I = [[([1]+[0]*(N-1)) if i == j else [0]*N for j in range(n)] for i in range(n)]
            OmX = [[[(u - v) % p for u, v in zip(I[i][j], X[i][j])] for j in range(n)] for i in range(n)]
            OpX = [[[(u + v) % p for u, v in zip(I[i][j], X[i][j])] for j in range(n)] for i in range(n)]
            try:
                g = mmul(inv_mat(OmX, N, p), OpX, N, p)
            except StopIteration:
                continue
            lhs = mmul(mmul(g, D, N, p), transpose(g), N, p)
            tested += 1
            if lhs != D:
                cay_ok = False; print('Cayley FAIL', p, a)
print('(b) Cayley transform lands in the stabilizer: %s (%d random tests)' % (cay_ok, tested))

# (c) defect formula
def dimO(nu):  # orbit dimension minus m(2m-1)N
    return -sum((4*(j+1)-3)*v for j, v in enumerate(nu))
def_ok = True
for m in range(1, 9):
    for k in range(1, 2*m):
        nus = [nu for nu in itertools.product([2, 1, 0], repeat=m) if sum(nu) == k and all(nu[i] >= nu[i+1] for i in range(m-1))]
        eps = k % 2
        top = max(nus, key=dimO)
        for nu in nus:
            n1 = nu.count(1)
            c = eps + dimO(top) - dimO(nu)
            d = n1*(n1+1) - c
            if d != n1:
                def_ok = False; print('defect mismatch', m, k, nu, d, n1)
print('(c) def(nu) = n_1(nu) for m <= 8, 1 <= k <= 2m-1: %s' % def_ok)
