"""
Function-level check for X = GL_4/Sp_4 over F = F_q((t)).
X(F) = nondegenerate alternating forms on F^4; K = GL_4(O) acts by x -> g x g^T.
K-orbits: x_mu = t^{a1} J_{12} + t^{a2} J_{34}, mu = (a1,a1,a2,a2), a1 >= a2.
(T * phi_0)(x) for the Hecke operator attached to lattices of a given relative type
= #{ lattices M with ed(M : O^4) = nu : x unimodular on M }.
We enumerate lattices tO^4 <= M <= t^{-1}O^4 by Hermite normal form and count,
for mu in {(2,2,0,0),(1,1,1,1)} (val Pf = 2) and nu with sum(nu) = -2, entries in {-1,0,1}.
A row is the function T_nu * phi_0 restricted to these two orbits (its full support,
see the bound in the report), so the 2x2 matrix decides whether 1_{O_(1,1,1,1)} is a
Hecke translate of phi_0 in degree 4.
"""
import itertools, sys
from fractions import Fraction

q = int(sys.argv[1]) if len(sys.argv) > 1 else 3

# Laurent polynomials over F_q as dict exp -> coeff (nonzero, mod q)
def add(a, b):
    c = dict(a)
    for e, v in b.items():
        c[e] = (c.get(e, 0) + v) % q
        if c[e] == 0:
            del c[e]
    return c
def mul(a, b):
    c = {}
    for e1, v1 in a.items():
        for e2, v2 in b.items():
            c[e1+e2] = (c.get(e1+e2, 0) + v1*v2) % q
    return {e: v for e, v in c.items() if v}
def scal(s, a):
    return {e: (s*v) % q for e, v in a.items() if (s*v) % q}
def shift(a, d):
    return {e+d: v for e, v in a.items()}
def integral(a):
    return all(e >= 0 for e in a)
def mono(e, c=1):
    return {e: c % q} if c % q else {}

def laurent_reps(lo, hi):
    """all polynomials sum_{k=lo}^{hi-1} c_k t^k (reps of t^lo O / t^hi O)"""
    exps = list(range(lo, hi))
    for coeffs in itertools.product(range(q), repeat=len(exps)):
        yield {e: c for e, c in zip(exps, coeffs) if c}

def lattices(d):
    """Lattices M = h O^4, h upper triangular, h_ii = t^{d_i}, h_ij (i<j) in t^{-1}O/t^{d_i}O,
    with t O^4 <= M <= t^{-1} O^4. Yields the column basis of h."""
    n = 4
    # columns: h e_j = sum_i h_ij e_i, i <= j
    offdiag_slots = [(i, j) for j in range(n) for i in range(j)]
    rep_lists = [list(laurent_reps(-1, d[i])) for (i, j) in offdiag_slots]
    for choice in itertools.product(*rep_lists):
        h = [[{} for _ in range(n)] for _ in range(n)]
        for i in range(n):
            h[i][i] = mono(d[i])
        for (i, j), r in zip(offdiag_slots, choice):
            h[i][j] = r
        # check t e_k in M for all k: solve h c = t e_k by back substitution, need c integral
        ok = True
        for k in range(n):
            rhs = [mono(1) if i == k else {} for i in range(n)]
            c = [{} for _ in range(n)]
            for j in range(n-1, -1, -1):
                # row j: sum_{l>=j} h[j][l] c[l] = rhs[j]  -> c[j] = (rhs[j] - sum_{l>j} h[j][l] c[l]) / t^{d_j}
                s = rhs[j]
                for l in range(j+1, n):
                    s = add(s, scal(-1, mul(h[j][l], c[l])))
                c[j] = shift(s, -d[j])
                if not integral(c[j]):
                    ok = False
                    break
            if not ok:
                break
        if ok:
            yield [[h[i][j] for i in range(n)] for j in range(n)]  # list of column vectors

def form(mu):
    n = 4
    X = [[{} for _ in range(n)] for _ in range(n)]
    X[0][1] = mono(mu[0]); X[1][0] = mono(mu[0], -1)
    X[2][3] = mono(mu[2]); X[3][2] = mono(mu[2], -1)
    return X

def pairing(X, u, v):
    s = {}
    for i in range(4):
        for j in range(4):
            if X[i][j] and u[i] and v[j]:
                s = add(s, mul(mul(u[i], X[i][j]), v[j]))
    return s

def unimodular(X, cols, mu, nu):
    # integrality on generators; determinant valuation = sum(mu) + 2 sum(nu) assumed 0
    for a in range(4):
        for b in range(a+1, 4):
            if not integral(pairing(X, cols[a], cols[b])):
                return False
    return True

def val(a):
    return min(a) if a else None
def det_poly(rows):
    n = len(rows)
    if n == 1:
        return rows[0][0]
    s = {}
    for j in range(n):
        if not rows[0][j]:
            continue
        minor = [[rows[i][k] for k in range(n) if k != j] for i in range(1, n)]
        term = mul(rows[0][j], det_poly(minor))
        s = add(s, scal((-1)**j, term))
    return s
def ed_type(cols):
    """elementary divisors of the lattice spanned by cols relative to O^4, sorted decreasing"""
    n = 4
    h = [[cols[j][i] for j in range(n)] for i in range(n)]
    e = [0]
    for k in range(1, n+1):
        best = None
        for R in itertools.combinations(range(n), k):
            for C in itertools.combinations(range(n), k):
                v = val(det_poly([[h[i][j] for j in C] for i in R]))
                if v is not None and (best is None or v < best):
                    best = v
        e.append(best)
    nu_asc = [e[k] - e[k-1] for k in range(1, n+1)]
    return tuple(sorted(nu_asc, reverse=True))

mus = [(2, 2, 0, 0), (1, 1, 1, 1)]
types = [(0, 0, -1, -1), (1, -1, -1, -1)]
counts = {(nu, mu): 0 for nu in types for mu in mus}
total = {nu: 0 for nu in types}
for d in set(itertools.permutations((0, 0, -1, -1))) | set(itertools.permutations((1, -1, -1, -1))):
    for cols in lattices(d):
        nu = ed_type(cols)
        if nu not in types:
            continue
        total[nu] += 1
        for mu in mus:
            if unimodular(form(mu), cols, mu, nu):
                counts[(nu, mu)] += 1
print("q =", q)
for nu in types:
    print("type", nu, "#lattices =", total[nu], " row (values on O_(2,2,0,0), O_(1,1,1,1)) =",
          [counts[(nu, mu)] for mu in mus])
A = [[counts[(nu, mu)] for mu in mus] for nu in types]
det = A[0][0]*A[1][1] - A[0][1]*A[1][0]
print("det =", det)
