# G2 audit: independent lattice counts for GL_{2m}/Sp_{2m} via Hermite normal forms of SUBlattices
# and duality (a different parametrisation from the run's scripts).  python3 -I, stdlib only.
#
# Convention of the audited nodes: for g in K t^lam K, x in g X(O) iff w_x is unimodular on
# M = g^{-T} O^n, and M has relative type -lam.  Hence
#     (T_lam * phi_0)(x) = #{ M of relative type -lam : w_x unimodular on M }.
# We produce the lattices M from sublattices L of O^n (column Hermite normal form B: upper
# triangular, B_ii = t^{d_i}, d_i >= 0, B_ij a polynomial of degree < d_i; every such B gives a
# distinct L <= O^n):
#   lam = (1,1,0,0) / (2,0,0,0) / omega_k :  M = L^*  (dual w.r.t. the dot product), L of type lam;
#   lam = (1,1,1,-1)                       :  M = t^{-1} L, L of type (2,0,0,0).
# Gram matrices: on L^* (basis B^{-T}) it is B^{-1} w B^{-T}; on t^{-1}L it is t^{-2} B^T w B.
# Unimodular = integral + unit Pfaffian; we test integrality and check v(det Gram) = 0.
import itertools

def lp_add(a, b, p):
    c = dict(a)
    for e, v in b.items():
        w = (c.get(e, 0) + v) % p
        if w: c[e] = w
        elif e in c: del c[e]
    return c

def lp_neg(a, p):
    return {e: (-v) % p for e, v in a.items()}

def lp_mul(a, b, p):
    c = {}
    for e1, v1 in a.items():
        for e2, v2 in b.items():
            e = e1 + e2
            w = (c.get(e, 0) + v1 * v2) % p
            if w: c[e] = w
            elif e in c: del c[e]
    return c

def lp_val(a):
    return min(a) if a else None

def mat_mul(A, B, p):
    n, k, m = len(A), len(B), len(B[0])
    C = [[{} for _ in range(m)] for _ in range(n)]
    for i in range(n):
        for j in range(m):
            s = {}
            for l in range(k):
                if A[i][l] and B[l][j]:
                    s = lp_add(s, lp_mul(A[i][l], B[l][j], p), p)
            C[i][j] = s
    return C

def transpose(A):
    return [list(r) for r in zip(*A)]

def scale(A, e, p):
    return [[lp_mul(x, {e: 1}, p) if x else {} for x in row] for row in A]

def det(A, p):
    n = len(A)
    if n == 1: return A[0][0]
    s = {}
    for j in range(n):
        if not A[0][j]: continue
        minor = [row[:j] + row[j+1:] for row in A[1:]]
        term = lp_mul(A[0][j], det(minor, p), p)
        s = lp_add(s, term if j % 2 == 0 else lp_neg(term, p), p)
    return s

def inverse_upper(B, p):
    n = len(B)
    d = [next(iter(B[i][i])) for i in range(n)]
    X = [[{} for _ in range(n)] for _ in range(n)]
    for j in range(n):
        X[j][j] = {-d[j]: 1}
        for i in range(j - 1, -1, -1):
            s = {}
            for k in range(i + 1, j + 1):
                if B[i][k] and X[k][j]:
                    s = lp_add(s, lp_mul(B[i][k], X[k][j], p), p)
            X[i][j] = lp_mul({-d[i]: p - 1}, s, p) if s else {}
    return X

def sublattices(n, total, p, dmax):
    """All L <= O^n with v(det) = total, d_i <= dmax, as column HNF matrices B."""
    for d in itertools.product(range(dmax + 1), repeat=n):
        if sum(d) != total: continue
        slots = [(i, j, d[i]) for i in range(n) for j in range(i + 1, n) if d[i] > 0]
        choices = [list(itertools.product(range(p), repeat=s[2])) for s in slots]
        for combo in itertools.product(*choices):
            B = [[{} for _ in range(n)] for _ in range(n)]
            for i in range(n): B[i][i] = {d[i]: 1}
            for (i, j, di), coeffs in zip(slots, combo):
                B[i][j] = {e: c for e, c in enumerate(coeffs) if c}
            yield B

def is_cyclic_colength2(B, p):
    # O^n / L of length 2: type (2,0,..) iff t O^n is NOT inside L, i.e. t B^{-1} not integral
    X = inverse_upper(B, p)
    return any(X[i][j] and lp_val(X[i][j]) + 1 < 0 for i in range(len(B)) for j in range(len(B)))

def killed_by_t(B, p):
    X = inverse_upper(B, p)
    return all((not X[i][j]) or lp_val(X[i][j]) + 1 >= 0 for i in range(len(B)) for j in range(len(B)))

def form_matrix(blocks, p):
    n = 2 * len(blocks)
    W = [[{} for _ in range(n)] for _ in range(n)]
    for i, e in enumerate(blocks):
        W[2*i][2*i+1] = {e: 1}
        W[2*i+1][2*i] = {e: p - 1}
    return W

def integral_unimodular(G, p):
    n = len(G)
    if any(G[i][j] and lp_val(G[i][j]) < 0 for i in range(n) for j in range(n)):
        return False
    assert lp_val(det(G, p)) == 0, "integral but not unimodular: index bookkeeping broken"
    return True

def gram_dual(B, W, p):
    Bi = inverse_upper(B, p)
    return mat_mul(mat_mul(Bi, W, p), transpose(Bi), p)

def gram_tinv(B, W, p):
    return scale(mat_mul(mat_mul(transpose(B), W, p), B, p), -2, p)

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

def run_gl4(p):
    L2 = list(sublattices(4, 2, p, 2))
    cyc = [B for B in L2 if is_cyclic_colength2(B, p)]
    non = [B for B in L2 if not is_cyclic_colength2(B, p)]
    forms = {'(1,1)': form_matrix((1, 1), p), '(2,0)': form_matrix((2, 0), p)}
    out = {}
    out[(1,1,0,0)] = ({k: sum(integral_unimodular(gram_dual(B, W, p), p) for B in non) for k, W in forms.items()}, len(non))
    out[(2,0,0,0)] = ({k: sum(integral_unimodular(gram_dual(B, W, p), p) for B in cyc) for k, W in forms.items()}, len(cyc))
    out[(1,1,1,-1)] = ({k: sum(integral_unimodular(gram_tinv(B, W, p), p) for B in cyc) for k, W in forms.items()}, len(cyc))
    print(f"== GL_4, q={p}: values of T_lam * phi_0 at x_(1,1)=tJ and x_(2,0)=t^2 J12 + J34 ==")
    for lam, (res, size) in out.items():
        print(f"  lam={lam}: |K t^lam K/K|={size}  values={res}")
    print(f"  predicted: (1,1,0,0): (1,1)->{(p+1)*(p*p+1)}, (2,0)->1 ; (2,0,0,0) and (1,1,1,-1): (1,1)->0, (2,0)->{p*(p+1)}")
    print(f"  cell sizes predicted: |Gr(2,4)|={(p**4-1)*(p**4-p)//((p**2-1)*(p**2-p))}, |P^3(O/t^2)|={p**3*(p**3+p**2+p+1)}")

def run_fibres(m, p, ks, entries=4):
    n = 2 * m
    print(f"== GL_{n} (m={m}), q={p}: #F_omega_k(w_nu) for sorted nu in [0,{entries-1}]^{m}, |nu|=k ==")
    for k in ks:
        Ls = [B for B in sublattices(n, k, p, 1) if killed_by_t(B, p)]
        line = []
        for nu in itertools.product(range(entries), repeat=m):
            if sum(nu) != k or list(nu) != sorted(nu, reverse=True): continue
            W = form_matrix(nu, p)
            cnt = sum(1 for B in Ls if integral_unimodular(gram_dual(B, W, p), p))
            pred = lagr(nu.count(1), p) if max(nu) <= 2 else 0
            flag = '' if cnt == pred else ' MISMATCH'
            line.append(f"{nu}:{cnt}/{pred}{flag}")
        print(f"  k={k} |Gr(k,{n})|={len(Ls)}: " + "  ".join(line))

if __name__ == '__main__':
    import sys
    for p in (2, 3):
        run_gl4(p)
    run_fibres(2, 3, [1, 2, 3])
    run_fibres(3, 2, [1, 2, 3, 4, 5])