# Independent check of the GL4/Sp4 counts in A1-4, over O = F_p[[t]] (p = 2 here; p = 3 for the
# cheap counts).  Value formula: (T_lam * phi0)(J') = #{lattices M : inv(O^4, M) = -w0 lam,
# J' unimodular on M}.  We work with N = t^2 M, so that t^3 O^4 <= N <= O^4 covers all M with
# relative position entries in [-2, 1].
import sys, itertools

p = int(sys.argv[1]) if len(sys.argv) > 1 else 2
PREC = 8  # truncate power series at t^PREC (all quantities below need < PREC)

def pmul(a, b):
    c = [0] * PREC
    for i, x in enumerate(a):
        if x:
            for j, y in enumerate(b):
                if i + j < PREC and y:
                    c[i + j] = (c[i + j] + x * y) % p
    return c
def padd(a, b): return [(x + y) % p for x, y in zip(a, b)]
def psub(a, b): return [(x - y) % p for x, y in zip(a, b)]
def mono(c, k):
    v = [0] * PREC
    if k < PREC: v[k] = c % p
    return v
ZERO = [0] * PREC
def val(a):
    for i, x in enumerate(a):
        if x: return i
    return PREC  # 'infinity' at this precision

# symplectic basis order (e1, f1, e2, f2): J(e1,f1)=J(e2,f2)=1
def form(Jdiag, x, y):
    # Jdiag = (s1, s2) power-series scalars: J'(e1,f1)=s1, J'(e2,f2)=s2
    s1, s2 = Jdiag
    t1 = psub(pmul(x[0], y[1]), pmul(x[1], y[0]))
    t2 = psub(pmul(x[2], y[3]), pmul(x[3], y[2]))
    return padd(pmul(s1, t1), pmul(s2, t2))

def elementary_divisors(cols):
    # cols: list of 4 column vectors (power series), lattice N = span; return sorted valuations
    # via d_k = min valuation of k x k minors (fine at this precision for N >= t^3 O^4)
    import itertools as it
    def det(m):
        n = len(m)
        if n == 1: return m[0][0]
        s = ZERO
        for j in range(n):
            minor = [row[:j] + row[j+1:] for row in m[1:]]
            term = pmul(m[0][j], det(minor))
            s = padd(s, term) if j % 2 == 0 else psub(s, term)
        return s
    rows = [[cols[j][i] for j in range(4)] for i in range(4)]
    d = [0]
    for k in range(1, 5):
        best = PREC
        for R in it.combinations(range(4), k):
            for C in it.combinations(range(4), k):
                v = val(det([[rows[r][c] for c in C] for r in R]))
                best = min(best, v)
        d.append(best)
    return tuple(sorted((d[k] - d[k - 1] for k in range(1, 5)), reverse=True))

def series_from_digits(digs):
    v = [0] * PREC
    for i, x in enumerate(digs): v[i] = x
    return v

def hnf_lattices(maxe, total):
    # upper-triangular HNF: column j has diagonal t^{e_j}, entries h_{ij} (i<j) reduced mod t^{e_i}
    for e in itertools.product(range(maxe + 1), repeat=4):
        if sum(e) != total: continue
        slots = [(i, j) for j in range(4) for i in range(j)]
        ranges = [itertools.product(range(p), repeat=e[i]) for (i, j) in slots]
        for choice in itertools.product(*[list(r) for r in ranges]):
            cols = [[ZERO[:] for _ in range(4)] for _ in range(4)]  # cols[j][i]
            for j in range(4): cols[j][j] = mono(1, e[j])
            for (i, j), digs in zip(slots, choice):
                cols[j][i] = series_from_digits(digs)
            yield e, cols

def contains_t3(e, cols):
    # does span(cols) contain t^3 e_k for all k?  Solve upper-triangular system column by column.
    for k in range(4):
        target = [ZERO[:] for _ in range(4)]; target[k] = mono(1, 3)
        coef = [None] * 4
        rem = [r[:] for r in target]
        for j in range(3, -1, -1):
            # coefficient of column j determined by row j (diagonal t^{e_j})
            r = rem[j]
            if val(r) < e[j]: return False
            c = r[e[j]:] + [0] * e[j]  # divide by t^{e_j}
            for i in range(4):
                rem[i] = psub(rem[i], pmul(c, cols[j][i]))
        if any(val(x) < PREC for x in rem): return False
    return True

def count_at(Jdiag, scale_total, required_val, maxe):
    # Count N with t^maxe O^4 <= N <= O^4, [O^4:N] = p^scale_total, J'(N,N) in t^required_val O
    tally = {}
    for e, cols in hnf_lattices(maxe, scale_total):
        if maxe == 3 and not contains_t3(e, cols): continue
        ok = True
        for a in range(4):
            for b in range(a + 1, 4):
                if val(form(Jdiag, cols[a], cols[b])) < required_val:
                    ok = False; break
            if not ok: break
        if not ok: continue
        typ = elementary_divisors(cols)
        tally[typ] = tally.get(typ, 0) + 1
    return tally

ONE = mono(1, 0); T1 = mono(1, 1); T2 = mono(1, 2)

def lam_from_type(typ):
    c = [x - 2 for x in typ]            # inv(O^4, M), dominant
    return tuple(-x for x in reversed(c))  # lambda = -w0 c

# Point of label (1,1,1,1): J' = t J.  t J unimodular on M = t^{-2}N  <=>  [O^4:N] = p^6, J'(N,N) in t^4 O.
tal = count_at((T1, T1), 6, 4, 3)
print('p = %d, point tJ (label (1,1,1,1)):' % p)
for typ, n in sorted(tal.items()):
    print('   N of type', typ, ' lambda =', lam_from_type(typ), ' count', n)
# Point of label (2,2,0,0): J'' = t^2 J_12 + J_34.
tal2 = count_at((T2, ONE), 6, 4, 3)
print('p = %d, point t^2 J12 + J34 (label (2,2,0,0)):' % p)
for typ, n in sorted(tal2.items()):
    print('   N of type', typ, ' lambda =', lam_from_type(typ), ' count', n)
print('expected: lambda=(1,1,0,0) at (1,1,1,1): (p+1)(p^2+1) =', (p+1)*(p*p+1),
      '; at (2,2,0,0): 1;  lambda=(2,0,0,0) at (2,2,0,0): p(p+1) =', p*(p+1), ', at (1,1,1,1): 0')