# Independent lattice counts for GL_4/Sp_4 (A4-9, A3-2, A3-9).  Residue field F_q, q prime.
# Forms: w1 = t*J (orbit (1,1,1,1)),  w2 = t^2*J12 + J34 (orbit (2,2,0,0)).
# R2 = F_q[t]/t^2 ; an element is (c0,c1).  Vectors in R2^4 are 4-tuples of such pairs.
import itertools, sys
def run(q):
    def add(a,b): return ((a[0]+b[0])%q,(a[1]+b[1])%q)
    def mul(a,b): return ((a[0]*b[0])%q,(a[0]*b[1]+a[1]*b[0])%q)
    def neg(a): return ((-a[0])%q,(-a[1])%q)
    def inv(a):  # a unit: a0 != 0
        i0=pow(a[0],q-2,q); return (i0, (-a[1]*i0*i0)%q)
    Z=(0,0); O=(1,0)
    def J_pairs(form):  # form as dict {(i,j): coefficient in R2}, antisymmetric
        return form
    w1={(0,1):(0,1),(2,3):(0,1)}                 # t*J
    w2={(0,1):(0,0),(2,3):(1,0)}                 # t^2 J12 (=0 mod t^2) + J34
    def omega(form,a,b):
        s=Z
        for (i,j),c in form.items():
            s=add(s, mul(c, add(mul(a[i],b[j]), neg(mul(a[j],b[i])))))
        return s
    def integral(form,gens):
        return all(omega(form,g,h)==Z for g in gens for h in gens)
    const=lambda v: tuple((x%q,0) for x in v)
    E=[const([1 if k==i else 0 for k in range(4)]) for i in range(4)]
    tE=[tuple((0,x) for x,_ in e) for e in E]
    # --- T_{(1,1,0,0)}: M = O^4 + t^{-1} P, P a 2-plane in F_q^4 ; Mbar = t O^4 + P
    planes=[]
    for pivots in itertools.combinations(range(4),2):
        free=[k for k in range(4) if k not in pivots]
        # RREF rows: pivot columns identity, entries in free columns to the right of pivot
        slots=[(r,k) for r,p in enumerate(pivots) for k in free if k>p]
        for vals in itertools.product(range(q),repeat=len(slots)):
            rows=[[0]*4 for _ in range(2)]
            for r,p in enumerate(pivots): rows[r][p]=1
            for (r,k),v in zip(slots,vals): rows[r][k]=v
            planes.append(rows)
    assert len(planes)==(q**2+1)*(q**2+q+1)
    resA={}
    for name,form in (('(1,1,1,1)',w1),('(2,2,0,0)',w2)):
        resA[name]=sum(1 for P in planes if integral(form, tE+[const(r) for r in P]))
    # --- T_{(1,1,1,-1)}: M = t^{-1} L, L = ker(c: R2^4 -> R2) ; Mbar = ker c, c in P^3(R2)
    units=[(a,b) for a in range(1,q) for b in range(q)]
    elems=[(a,b) for a in range(q) for b in range(q)]
    cs=[]
    for c in itertools.product(elems,repeat=4):
        i0=next((i for i in range(4) if c[i][0]!=0),None)
        if i0 is None: continue
        if c[i0]!=O: continue        # normalise: first unit coordinate = 1
        cs.append((c,i0))
    assert len(cs)==q**3*(q**4-1)//(q-1)
    resB={}
    for name,form in (('(1,1,1,1)',w1),('(2,2,0,0)',w2)):
        cnt=0
        for c,i0 in cs:
            gens=[]
            for j in range(4):
                if j==i0: continue
                g=[Z]*4; g[j]=O; g[i0]=neg(c[j]); gens.append(tuple(g))
            if integral(form,gens): cnt+=1
        resB[name]=cnt
    # --- T_{(2,0,0,0)}: M = O^4 + O t^{-2} v, v primitive mod t^2 up to units; integral iff omega(v, e_i)=0 in R2 for all i
    resC={}
    for name,form in (('(1,1,1,1)',w1),('(2,2,0,0)',w2)):
        cnt=0
        for v,i0 in cs:   # same parametrisation of P^3(R2)
            if all(omega(form,v,e)==Z for e in E): cnt+=1
        resC[name]=cnt
    print('q=%d'%q)
    print('  T_(1,1,0,0)*phi0 :',resA,' total planes',len(planes),' predicted',{'(1,1,1,1)':(q+1)*(q*q+1),'(2,2,0,0)':1})
    print('  T_(1,1,1,-1)*phi0:',resB,' total',len(cs),' predicted',{'(1,1,1,1)':0,'(2,2,0,0)':q*q+q})
    print('  T_(2,0,0,0)*phi0 :',resC,' total',len(cs),' predicted',{'(1,1,1,1)':0,'(2,2,0,0)':q*q+q})
for q in (2,3): run(q)
