"""
Independent count of Hecke translates T_nu * phi_0 for X = GL_{2n}/Sp_{2n} over F_q((t)).

Method (different from the run's Hermite-normal-form enumeration):
  x = sum_p t^{a_p} J_{(2p-1,2p)} (symplectic elementary divisors (a_1,a_1,...,a_n,a_n)).
  (T_nu * phi_0)(x) = #{lattices M of relative type nu to O^{2n} : x unimodular on M}.
  x unimodular on M  <=>  M = M^* (x-dual).  For M in the box t^N O <= M <= t^{-N} O,
  self-duality forces A <= M <= A^*, A = t^N O + (t^{-N} O)^*, and M <-> L = M/A is a
  t-stable Lagrangian subspace of Q = A^*/A for the F_q-valued form omega = Res x(.,.).
  We enumerate all Lagrangians of (Q, omega) by isotropic RREF backtracking, keep the
  t-stable ones, and read off the relative type of M from the O-module M / t^N O^{2n}.
Usage: python3 -I selfdual_counts.py q
"""
import itertools, sys
from collections import Counter

def build(q, a_list, N):
    n2 = 2 * len(a_list)
    a_of = []
    for a in a_list:
        a_of += [a, a]
    basisQ = []     # (i, k) with lo_i <= k < hi_i : the element t^k e_i of Q
    lo, hi = [], []
    for i in range(n2):
        a = a_of[i]
        l = max(-N, -N - a)
        h = min(N, N - a)
        lo.append(l); hi.append(h)
        if h < l:
            return None
        for k in range(l, h):
            basisQ.append((i, k))
    idx = {b: j for j, b in enumerate(basisQ)}
    dimQ = len(basisQ)
    # omega matrix: omega(t^k e_i, t^l e_j) = Res(t^{k+l} x(e_i,e_j))
    W = [[0] * dimQ for _ in range(dimQ)]
    for u, (i, k) in enumerate(basisQ):
        for v, (j, l) in enumerate(basisQ):
            p_i, p_j = i // 2, j // 2
            if p_i != p_j or i == j:
                continue
            a = a_of[i]
            sign = 1 if i % 2 == 0 else -1     # x(e_{2p-1}, e_{2p}) = +t^a, reversed = -t^a
            if k + l + a == -1:
                W[u][v] = sign % q
    # t-action on Q: t^k e_i -> t^{k+1} e_i (zero if k+1 = hi_i)
    tmap = []
    for (i, k) in basisQ:
        tmap.append(idx.get((i, k + 1)))
    return dict(n2=n2, lo=lo, hi=hi, basisQ=basisQ, idx=idx, dimQ=dimQ, W=W, tmap=tmap, N=N)

def omega(W, u, v, q):
    s = 0
    for i, ui in enumerate(u):
        if ui:
            row = W[i]
            for j, vj in enumerate(v):
                if vj and row[j]:
                    s += ui * row[j] * vj
    return s % q

def lagrangians(q, dimQ, W):
    D = dimQ // 2
    for pivots in itertools.combinations(range(dimQ), D):
        pset = set(pivots)
        free = [c for c in range(dimQ) if c not in pset]
        row_free = [[c for c in free if c > p] for p in pivots]
        rows = []
        def rec(r):
            if r == D:
                yield [list(x) for x in rows]
                return
            for vals in itertools.product(range(q), repeat=len(row_free[r])):
                v = [0] * dimQ
                v[pivots[r]] = 1
                for c, x in zip(row_free[r], vals):
                    v[c] = x
                if all(omega(W, v, w, q) == 0 for w in rows):
                    rows.append(v)
                    yield from rec(r + 1)
                    rows.pop()
        yield from rec(0)

def rank_mod(vecs, q):
    rows = [list(v) for v in vecs if any(v)]
    r = 0
    if not rows:
        return 0
    ncol = len(rows[0])
    for c in range(ncol):
        piv = None
        for i in range(r, len(rows)):
            if rows[i][c] % q:
                piv = i; break
        if piv is None:
            continue
        rows[r], rows[piv] = rows[piv], rows[r]
        inv = pow(rows[r][c], q - 2, q)
        rows[r] = [(x * inv) % q for x in rows[r]]
        for i in range(len(rows)):
            if i != r and rows[i][c] % q:
                f = rows[i][c]
                rows[i] = [(x - f * y) % q for x, y in zip(rows[i], rows[r])]
        r += 1
        if r == len(rows):
            break
    return r

def t_stable(L, S, q):
    # L list of vectors in Q coordinates; check t(L) subset of span(L)
    r0 = rank_mod(L, q)
    for v in L:
        tv = [0] * S['dimQ']
        for j, x in enumerate(v):
            if x and S['tmap'][j] is not None:
                tv[S['tmap'][j]] = (tv[S['tmap'][j]] + x) % q
        if any(tv) and rank_mod(L + [tv], q) != r0:
            return False
    return True

def lattice_type(L, S, q):
    """relative type nu of M = A + lift(L) w.r.t. O^{2n}; via M / t^N O^{2n} inside
    V_big = t^{-N}O^{2n} / t^N O^{2n} with basis t^k e_i, -N <= k < N."""
    N, n2 = S['N'], S['n2']
    big = [(i, k) for i in range(n2) for k in range(-N, N)]
    bidx = {b: j for j, b in enumerate(big)}
    gens = []
    for v in L:
        w = [0] * len(big)
        for j, x in enumerate(v):
            if x:
                w[bidx[S['basisQ'][j]]] = x
        gens.append(w)
    for i in range(n2):
        for k in range(S['hi'][i], N):
            w = [0] * len(big); w[bidx[(i, k)]] = 1
            gens.append(w)
    def tshift(w):
        out = [0] * len(big)
        for j, x in enumerate(w):
            if x:
                i, k = big[j]
                if k + 1 < N:
                    out[bidx[(i, k + 1)]] = x
        return out
    dims = []
    cur = gens
    for kk in range(2 * N + 1):
        dims.append(rank_mod(cur, q))
        cur = [tshift(w) for w in cur]
    # #{l_i >= k} = dims[k-1] - dims[k]
    lengths = []
    for k in range(1, 2 * N + 1):
        cnt_ge_k = dims[k - 1] - dims[k]
        cnt_ge_k1 = (dims[k] - dims[k + 1]) if k + 1 <= 2 * N else 0
        lengths += [k] * (cnt_ge_k - cnt_ge_k1)
    lengths += [0] * (n2 - len(lengths))
    nu = sorted([N - l for l in lengths], reverse=True)
    return tuple(nu)

def count(q, a_list, N=1):
    S = build(q, a_list, N)
    if S is None:
        return Counter()
    if S['dimQ'] == 0:
        return Counter({lattice_type([], S, q): 1})
    C = Counter()
    for L in lagrangians(q, S['dimQ'], S['W']):
        if t_stable(L, S, q):
            C[lattice_type(L, S, q)] += 1
    return C

if __name__ == '__main__':
    q = int(sys.argv[1]) if len(sys.argv) > 1 else 3
    which = sys.argv[2] if len(sys.argv) > 2 else 'both'
    if which in ('gl4', 'both'):
        print('GL4/Sp4, q=%d, box tO <= M <= t^-1 O' % q)
        for mu in [(2, 0), (1, 1)]:
            C = count(q, list(mu))
            print('  orbit', (mu[0], mu[0], mu[1], mu[1]), dict(C))
        print('  A4-9 claims: R_A=(0,0,-1,-1) -> (1, (q+1)(q^2+1)) =', (1, (q + 1) * (q * q + 1)),
              '; R_B=(1,-1,-1,-1) -> (q^2+q, 0) =', (q * q + q, 0))
    if which in ('gl6', 'both'):
        print('GL6/Sp6, q=%d, box tO <= M <= t^-1 O' % q)
        for mu in [(2, 0, 0), (1, 1, 0), (2, 1, -1), (2, 2, -2)]:
            C = count(q, list(mu))
            print('  orbit', tuple(x for x in mu for _ in (0, 1)), dict(C))
        print('  A4-10 claims (columns c1..c4):')
        print('    R1=(0,0,0,0,-1,-1):', (1, (q + 1) * (q * q + 1), 0, 0))
        print('    R2=(1,0,0,-1,-1,-1):', ((q ** 4 - 1) // (q - 1), (q + 1) * (q * q + 1) * (q * q + q), (q + 1) ** 2, 0))
        print('    R3=(1,1,-1,-1,-1,-1):', ((q + 1) * (q * q + 1), 0, 0, 1))
        print('  G3 closed forms:')
        print('    R1:', (1, (q + 1) * (q * q + 1), 0, 0))
        print('    R2:', (q * (q + 1) * (q * q + 1), q ** 3 * (q + 1) ** 2 * (q * q + 1), (q + 1) ** 2, 0))
        print('    R3:', (q ** 3 * (q + 1) * (q * q + 1), 0, q ** 3 - q, 1))