# Finite check for X = A^n, n=2, over F_p: orbits of (A/J)^x on Y_{N,M} = (t^-N O / t^M O)^n.
# Elements of O/t^K represented as tuples of K coefficients mod p (truncated power series).
import itertools
p=2; n=2
def mul(a,b,K):
    c=[0]*K
    for i in range(K):
        if a[i]:
            for j in range(K-i):
                c[i+j]=(c[i+j]+a[i]*b[j])%p
    return tuple(c)
def add(a,b): return tuple((x+y)%p for x,y in zip(a,b))
def orbits(N,M):
    K=N+M  # J = t^{N+M} M_n(O); vectors t^-N O/t^M O  ~ shift by t^N: O/t^{N+M}
    series=list(itertools.product(range(p),repeat=K))
    # Iwahori order mod t^K: upper-tri mod t: entry (2,1) in tO
    group=[]
    for a11 in series:
        if a11[0]==0: continue
        for a22 in series:
            if a22[0]==0: continue
            for a12 in series:
                for a21 in series:
                    if a21[0]!=0: continue
                    group.append(((a11,a12),(a21,a22)))
    vecs=list(itertools.product(series,repeat=n))
    def act(g,v):
        return tuple(add(mul(g[i][0],v[0],K),mul(g[i][1],v[1],K)) for i in range(n))
    seen=set(); orbs=[]
    for v in vecs:
        if v in seen: continue
        orb={act(g,v) for g in group}
        seen|=orb; orbs.append((v,len(orb)))
    return orbs
for (N,M) in [(0,1),(0,2),(1,1),(1,2)]:
    o=orbits(N,M)
    print("N,M=",(N,M),"#orbits=",len(o),"expected",(N+M)*n+1,"sizes",sorted(s for _,s in o))
# omega e1 = t e_n, omega e_i = e_{i-1}: check omega L_j = L_{j-1} via generators (valuation vectors)
def L(j): return [0]*j+[1]*(n-j)   # L_j = t^{(0^j,1^{n-j})} O^n, as exponent vector per coordinate
def omega_exp(ex):  # omega maps coordinate e_i to e_{i-1} (i>=2), e_1 to t e_n; lattice t^ex O^n -> exponents
    new=[None]*n
    for i in range(n):       # basis vector t^{ex[i]} e_{i+1} (0-indexed i)
        if i==0: new[n-1]=ex[0]+1
        else: new[i-1]=ex[i]
    return new
for j in range(1,n+1):
    print("omega L_%d ="%j, omega_exp(L(j)), " L_%d ="%(j-1), L(j-1))