# Independent check of A5-6's m=3 table and A5-5's N((1,0,0,0),(1,0)).  O = F_q[[t]], q prime.
# N(lambda,b) = #{Lambda'' >= Lambda_0 of type lambda : x_b unimodular on Lambda''}, x_b = sum t^{b_i} J_2 (blocks (2i-1,2i)).
# Types handled: lambda=(2,2,0^{2m-2}) [Lambda''/Lambda_0 = (O/t^2)^2], lambda=(1^k,0^{2m-k}) [= k-plane P, Lambda''=Lambda_0+t^{-1}P].
import itertools, sys
def poly_mul(a,b,q,N):
    c=[0]*N
    for i,x in enumerate(a):
        if x:
            for j,y in enumerate(b):
                if i+j<N: c[i+j]=(c[i+j]+x*y)%q
    return c
def form(b,q,N,u,v):   # x_b(u,v) mod t^N ; u,v lists of polys (length N) per coordinate
    m=len(b); s=[0]*N
    for i in range(m):
        if b[i]>=N: continue
        tb=[0]*N; tb[b[i]]=1
        a=poly_mul(u[2*i],v[2*i+1],q,N); c=poly_mul(u[2*i+1],v[2*i],q,N)
        d=[(x-y)%q for x,y in zip(a,c)]
        d=poly_mul(tb,d,q,N)
        s=[(x+y)%q for x,y in zip(s,d)]
    return s
def rref_subspaces(k,n,q):
    for pivots in itertools.combinations(range(n),k):
        free=[c for c in range(n) if c not in pivots]
        slots=[(r,c) for r,p in enumerate(pivots) for c in free if c>p]
        for vals in itertools.product(range(q),repeat=len(slots)):
            rows=[[0]*n for _ in range(k)]
            for r,p in enumerate(pivots): rows[r][p]=1
            for (r,c),v in zip(slots,vals): rows[r][c]=v
            yield pivots,rows
def count_planes_type11(b,q):   # lambda=(1^k,0..): Lambda''=Lambda_0+t^{-1}P; integral iff x(P,Lambda_0) in tO and x(P,P) in t^2 O
    m=len(b); n=2*m; k=sum(b)   # need |lambda|=|b| for unimodularity; we use k=|b|
    E=[[[1 if (c==j and d==0) else 0 for d in range(2)] for c in range(n)] for j in range(n)]
    cnt=0
    for piv,rows in rref_subspaces(k,n,q):
        gens=[[[r[c],0] for c in range(n)] for r in rows]
        ok=all(form(b,q,2,g,e)[0]==0 for g in gens for e in E) and all(form(b,q,2,g,h)==[0,0] for g in gens for h in gens)
        if ok: cnt+=1
    return cnt
def count_type22(b,q):          # lambda=(2,2,0..): Lambda''=Lambda_0+t^{-2}(O u+O v), u,v canonical lifts of an RREF plane with t-parts on non-pivot columns
    m=len(b); n=2*m
    E=[[[1 if (c==j and d==0) else 0 for d in range(4)] for c in range(n)] for j in range(n)]
    cnt=0
    for piv,rows in rref_subspaces(2,n,q):
        # prune: mod t condition x_b(rowbar, e)=0 mod t
        g0=[[[r[c],0,0,0] for c in range(n)] for r in rows]
        if not all(form(b,q,1,g,e)[0]==0 for g in g0 for e in E): continue
        free=[c for c in range(n) if c not in piv]
        for w1 in itertools.product(range(q),repeat=len(free)):
            u=[[rows[0][c],0,0,0] for c in range(n)]
            for c,x in zip(free,w1): u[c][1]=x
            if not all(form(b,q,2,u,e)[:2]==[0,0] for e in E): continue
            for w2 in itertools.product(range(q),repeat=len(free)):
                v=[[rows[1][c],0,0,0] for c in range(n)]
                for c,x in zip(free,w2): v[c][1]=x
                if not all(form(b,q,2,v,e)[:2]==[0,0] for e in E): continue
                if form(b,q,4,u,v)==[0,0,0,0]: cnt+=1
    return cnt
q=int(sys.argv[1])
print('q=%d'%q)
print('A5-5 check N((1,0,0,0),(1,0)) =',count_planes_type11((1,0),q),' predicted',q+1)
cols=[(4,0,0),(3,1,0),(2,2,0),(2,1,1)]
pred22={(4,0,0):1,(3,1,0):q*(q*q-1),(2,2,0):q**3*(q+1)*(q*q+1),(2,1,1):0}
pred1111={(4,0,0):0,(3,1,0):0,(2,2,0):1,(2,1,1):(q+1)*(q*q+1)}
for b in cols:
    print(' row (1,1,1,1,0,0), col',b,': N =',count_planes_type11(b,q),' predicted',pred1111[b])
for b in cols:
    if q==3 and b==(2,1,1): print(' row (2,2,0,0,0,0), col',b,': skipped at q=3 (72M lifts); predicted 0 by (V)'); continue
    print(' row (2,2,0,0,0,0), col',b,': N =',count_type22(b,q),' predicted',pred22[b])