# A9: algebraic checks on the orbit-indicator module of GL_2/T (formulas of A6-2), q an integer.
from fractions import Fraction as Fr
import sys
MM=30
def basis():
    B=['A+','A-']+[(i,m) for m in range(MM+1) for i in '01']; return B
def sw(X): return X if X in ('A+','A-') else (('1' if X[0]=='0' else '0'),X[1])
def T(s,O,q):
    if s=='s0': return {sw(k):v for k,v in T('s1',sw(O),q).items()}
    if O in ('A+','A-'): return {('A-' if O=='A+' else 'A+'):1,('0',0):1}
    i,m=O
    if O==('0',0): return {'A+':q-1,'A-':q-1,('0',0):q-2}
    if i=='1': return {('0',m+1):1}
    return {('1',m-1):q,('0',m):q-1}
def Tw(O): return {('A-' if O=='A+' else 'A+') if O in ('A+','A-') else sw(O):1}
def app(op,vec,q=None):
    out={}
    for O,c in vec.items():
        img=op(O) if q is None else op(O,q)
        for k,v in img.items(): out[k]=out.get(k,0)+c*v
    return {k:v for k,v in out.items() if v}
def lam(vec): return sum(c*({'A+':1,'A-':-1}.get(O,0)) for O,c in vec.items())
def psi(vec): return sum(c*(Fr(1,2) if O in ('A+','A-') else (-1)**(O[1]+1)) for O,c in vec.items())
bad=0
for q in (2,3,4,5,7,9):
    for O in basis():
        if O not in ('A+','A-') and O[1]>MM-2: continue
        for s in ('s0','s1'):
            v=app(lambda X,qq: T(s,X,qq),{O:1},q)
            if lam(v)!=-lam({O:1}) or psi(v)!=-psi({O:1}): bad+=1; print('functional',q,s,O)
            v2=app(lambda X,qq: T(s,X,qq),v,q)   # (T-q)(T+1)=T^2-(q-1)T-q
            chk={k:v2.get(k,0)-(q-1)*v.get(k,0)-(q if k==O else 0) for k in set(v2)|set(v)|{O}}
            if any(chk.values()): bad+=1; print('quadratic',q,s,O)
        w=app(Tw,{O:1})
        if lam(w)!=-lam({O:1}) or psi(w)!=psi({O:1}): bad+=1; print('omega',O)
        # T_w T_s0 T_w = T_s1  (T_w^2 = 1 on this module)
        a=app(Tw,app(lambda X,qq: T('s0',X,qq),app(Tw,{O:1}),q)); b=app(lambda X,qq: T('s1',X,qq),{O:1},q)
        if a!=b: bad+=1; print('conj',q,O)
    # f_n and f^(1)_n, their functionals, and the codimension of their span
    def fn(n,i):
        out={}
        # f_n(gamma)=max(0,n+1-d(v_i,gamma)); d(v0,.) on classes: A:0, (0,m):m, (1,m):m+1
        for O in basis():
            if O in ('A+','A-'): d=0
            else:
                d=O[1] if O[0]==i else O[1]+1
            val=max(0,n+1-d)
            if val: out[O]=val
        return out
    N=12
    vecs=[fn(n,i) for n in range(N+1) for i in '01']
    for v in vecs:
        if lam(v)!=0 or psi(v)!=0: bad+=1; print('f_n not killed',q)
    # Hecke-stability: (T_s0+1)f_n = f1_{n+1}+q f1_{n-1}
    for n in range(N):
        lhs=app(lambda X,qq: T('s0',X,qq),fn(n,'0'),q); lhs={k:lhs.get(k,0)+fn(n,'0').get(k,0) for k in set(lhs)|set(fn(n,'0'))}
        rhs=fn(n+1,'1'); 
        if n>=1:
            for k,v in fn(n-1,'1').items(): rhs[k]=rhs.get(k,0)+q*v
        if {k:v for k,v in lhs.items() if v}!={k:v for k,v in rhs.items() if v}: bad+=1; print('shadow',q,n)
    # rank of span{f_n,f1_n : n<=N} inside span{1_O : O in Y_N u omega Y_N}
    cols=['A+','A-']+[(i,m) for m in range(N+1) for i in '01']
    rows=[[Fr(v.get(c,0)) for c in cols] for v in vecs]
    # Gaussian elimination
    r=0; Mx=[row[:] for row in rows]
    for c in range(len(cols)):
        piv=next((k for k in range(r,len(Mx)) if Mx[k][c]!=0),None)
        if piv is None: continue
        Mx[r],Mx[piv]=Mx[piv],Mx[r]
        for k in range(len(Mx)):
            if k!=r and Mx[k][c]!=0:
                f=Mx[k][c]/Mx[r][c]; Mx[k]=[x-f*y for x,y in zip(Mx[k],Mx[r])]
        r+=1
    if not (r==len(vecs) and len(cols)-r==2): bad+=1; print('rank',q,r,len(cols))
    # invariant measure consistency: c_{O'O} mu(O') = c_{OO'} mu(O), mu(A)=1, mu(i,m)=(q-1)q^m
    mu=lambda O: 1 if O in ('A+','A-') else (q-1)*q**O[1]
    for s in ('s0','s1'):
        for O in basis():
            if O not in ('A+','A-') and O[1]>MM-2: continue
            for Op,c in T(s,O,q).items():
                if c*mu(Op)!=T(s,Op,q).get(O,0)*mu(O): bad+=1; print('measure',q,s,O,Op)
print('module checks done, failures:',bad)