"""Hecke counts for X = GL_{2m}/Sp_{2m} through Lagrangian submodules.

N(lam, c) := (T_lam * phi_0)(x_c) = #{Lambda'' containing Lambda_0 = O^{2m}, of type lam,
on which x_c = (+) t^{c_i} J_2 is unimodular}.  A lattice Lambda'' containing Lambda_0 is
self-dual for x_c iff Lambda''/Lambda_0 is a Lagrangian submodule of the finite symplectic
module M'_c = Lambda_0^#/Lambda_0 = (+)_i (O/t^{c_i})^2 (F/O-valued form).  So
N(lam, c) = #{Lagrangian O-submodules of M'_c of module type lam}.

O = F_p[[t]], p prime.  M'_c is an F_p-space with nilpotent t; O-submodules = t-stable
subspaces; Lagrangian for the F/O form = Lagrangian for its residue (coefficient of 1/t).

Usage: python3 -I symp_lagr.py
"""
import sys, os, itertools
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fplin import rref, nullspace


def build(c):
    """Index the F_p-basis t^j g_{i,s}, j in [-c_i,-1]; return (index list, t-matrix, B)."""
    idx = []
    for i, ci in enumerate(c):
        for s in (0, 1):
            for j in range(-ci, 0):
                idx.append((i, s, j))
    pos = {k: n for n, k in enumerate(idx)}
    N = len(idx)
    # t action: column images
    tmap = [None] * N
    for n, (i, s, j) in enumerate(idx):
        tmap[n] = pos.get((i, s, j + 1))  # None means 0
    B = [[0] * N for _ in range(N)]
    for n, (i, s, j) in enumerate(idx):
        for n2, (i2, s2, j2) in enumerate(idx):
            if i == i2 and j + j2 == -1 - c[i]:
                if s == 0 and s2 == 1:
                    B[n][n2] = 1
                elif s == 1 and s2 == 0:
                    B[n][n2] = -1
    return idx, tmap, B


def tvec(v, tmap, p):
    w = [0] * len(v)
    for n, x in enumerate(v):
        if x and tmap[n] is not None:
            w[tmap[n]] = (w[tmap[n]] + x) % p
    return tuple(w)


def cyclic(v, tmap, p):
    out = []
    while any(v):
        out.append(v)
        v = tvec(v, tmap, p)
    return out


def module_type(W, tmap, p):
    """Partition type of the t-stable subspace W (RREF basis)."""
    dims = []
    cur = list(W)
    while cur:
        R = rref(cur, p)
        dims.append(len(R))
        cur = [tvec(v, tmap, p) for v in R]
        cur = [v for v in cur if any(v)]
    dims.append(0)
    # number of parts >= k+1 is dims[k]-dims[k+1]
    parts = []
    for k in range(len(dims) - 1):
        cnt = dims[k] - dims[k + 1]
        parts.append(cnt)
    # parts[k] = #parts of size exactly... convert: #parts>=k+1 = dims[k]-dims[k+1]? no:
    # dim t^k W = sum_i max(lam_i - k, 0); so dims[k]-dims[k+1] = #{i: lam_i >= k+1}
    lam = []
    ge = parts  # ge[k] = #{i : lam_i >= k+1}
    for size in range(len(ge), 0, -1):
        n_exact = ge[size - 1] - (ge[size] if size < len(ge) else 0)
        lam += [size] * n_exact
    return tuple(lam)


def perp(W, B, N, p):
    rows = []
    for w in W:
        rows.append([sum(w[a] * B[a][b] for a in range(N)) % p for b in range(N)])
    return nullspace(rows, N, p) if rows else [tuple(1 if k == n else 0 for k in range(N)) for n in range(N)]


def lagrangians(c, p):
    idx, tmap, B = build(c)
    N = len(idx)
    half = N // 2
    levels = {0: {tuple()}}
    for d in range(0, half):
        for W in levels.get(d, ()):
            P = perp(W, B, N, p)
            P = rref(list(P), p)
            # complement of W inside P: pivot-free choice -- just iterate all vectors of P
            # but skip those already in W by checking dim growth.
            for coeffs in itertools.product(range(p), repeat=len(P)):
                if not any(coeffs):
                    continue
                # normalise: first nonzero coefficient 1 (scalars give same module)
                fnz = next(x for x in coeffs if x)
                if fnz != 1:
                    continue
                v = tuple(sum(cf * P[k][n] for k, cf in enumerate(coeffs)) % p for n in range(N))
                W2 = rref(list(W) + cyclic(v, tmap, p), p)
                if len(W2) == len(W):
                    continue
                levels.setdefault(len(W2), set()).add(W2)
    return levels.get(half, set()), tmap, B, idx


def check_isotropic(W, B, p):
    for a in W:
        for b in W:
            if sum(a[i] * B[i][j] * b[j] for i in range(len(a)) for j in range(len(b))) % p:
                return False
    return True


def counts(c, p):
    L, tmap, B, idx = lagrangians(c, p)
    res = {}
    for W in L:
        assert check_isotropic(W, B, p)
        ty = module_type(W, tmap, p)
        res[ty] = res.get(ty, 0) + 1
    return res


def dom_ge(a, b):
    """a >=_dom b for partitions of the same size (padded)."""
    n = max(len(a), len(b))
    a = list(a) + [0] * (n - len(a)); b = list(b) + [0] * (n - len(b))
    sa = sb = 0
    for x, y in zip(a, b):
        sa += x; sb += y
        if sa < sb:
            return False
    return sa == sb


if __name__ == '__main__':
    for p in (2, 3):
        print('=== p =', p)
        print('GL4/Sp4 columns b=(2,0),(1,1),(1,0):')
        for c in [(2, 0), (1, 1), (1, 0)]:
            print('  b =', c, counts(c, p))
        print('GL6/Sp6 columns b=(4,0,0),(3,1,0),(2,2,0),(2,1,1):')
        for c in [(4, 0, 0), (3, 1, 0), (2, 2, 0), (2, 1, 1)]:
            r = counts(c, p)
            print('  b =', c, dict(sorted(r.items(), reverse=True)))