# Affine Lusztig-Vogan module of GL2/T (= PGL2/T) in the tree model.
# Orbits: 'P' = A_+, 'M' = A_-, (i,n) n>=0 : closer vertex u_i at distance n, other at n+1; (i,0)=B_i.
# Laurent polys in v as dict exp->int.  q = v^{-2}.  Standard basis M_O = v^{d_O} 1_O.
import sys
from collections import defaultdict

def padd(a,b,s=1):
    r=dict(a)
    for e,c in b.items():
        r[e]=r.get(e,0)+s*c
        if r[e]==0: del r[e]
    return r
def pmul(a,b):
    r={}
    for e1,c1 in a.items():
        for e2,c2 in b.items():
            r[e1+e2]=r.get(e1+e2,0)+c1*c2
    return {e:c for e,c in r.items() if c}
def mono(e,c=1): return {e:c} if c else {}
Q=mono(-2)
ONE=mono(0)
def d(O):
    return 0 if O in ('P','M') else O[1]+1

# module elements: dict orbit -> poly, in the 1_O basis ("functions")
def eadd(x,y,s=1):
    r=dict(x)
    for O,p in y.items():
        r[O]=padd(r.get(O,{}),p,s)
        if not r[O]: del r[O]
    return r
def escale(x,p):
    r={}
    for O,a in x.items():
        b=pmul(a,p)
        if b: r[O]=b
    return r

def T_on_basis(s,O):
    """T_s 1_O as dict orbit->poly (q generic). s in {0,1}."""
    qm1=padd(Q,ONE,-1); qm2=padd(Q,mono(0,2),-1)
    j = 1 if s==1 else 0   # s1 fixes u_0 (vertex type 0): block index n=d(u_0); s0 fixes u_1
    # for s1: zero block {P,M,(0,0)}, higher blocks {(1,n-1),(0,n)}; for s0 swap roles 0<->1
    a = 0 if s==1 else 1   # the 'fixed vertex type'
    b = 1-a
    if O in ('P','M'):
        other = 'M' if O=='P' else 'P'
        return {other:ONE, (a,0):ONE}
    i,n=O
    if (i,n)==(a,0):
        return {'P':qm1,'M':qm1,(a,0):qm2}
    if i==b:   # (b,n-1) -> (a,n)
        return {(a,n+1):ONE}
    if i==a:   # (a,n), n>=1 -> q (b,n-1) + (q-1)(a,n)
        return {(b,n-1):Q,(a,n):qm1}
def T(s,x):
    r={}
    for O,p in x.items():
        r=eadd(r,escale(T_on_basis(s,O),p))
    return r
def Tw(x):
    r={}
    for O,p in x.items():
        O2 = {'P':'M','M':'P'}[O] if O in ('P','M') else (1-O[0],O[1])
        r=eadd(r,{O2:p})
    return r
def Cs(s,x):  # underline H_s = v(T_s+1)
    return escale(eadd(T(s,x),x),mono(1))
# conversion between 1_O basis and M_O basis
def toM(x): return {O:pmul(p,mono(-d(O))) for O,p in x.items()}
def fromM(y): return {O:pmul(p,mono(d(O))) for O,p in y.items()}

def M(O): return fromM({O:ONE})

# bar involution on M-basis, computed from bar(M_P)=M_P, bar(h m)=bar(h)bar(m)
def pbar(p): return {-e:c for e,c in p.items()}
barcache={}
def Hs(s,x): return escale(T(s,x),mono(1))  # H_s = v T_s
def Hsbar(s,x): # bar(H_s) = H_s + v - v^{-1}
    return eadd(Hs(s,x), escale(x,padd(mono(1),mono(-1),-1)))
def bar_basis(O):
    if O in barcache: return barcache[O]
    if O in ('P','M'):
        r=M(O)
    elif O==(0,0):
        # M_{(0,0)} = H_{s1} M_P - v T_w M_P
        r=eadd(Hsbar(1,M('P')), escale(Tw(M('P')),mono(-1)),-1)
    elif O==(1,0):
        r=Tw(bar_basis((0,0)))
    else:
        i,n=O
        # M_{(1,n)} = H_{s0} M_{(0,n-1)}, M_{(0,n)} = H_{s1} M_{(1,n-1)}
        s = 0 if i==1 else 1
        r=Hsbar(s,bar_basis((1-i,n-1)))
    barcache[O]=r
    return r
def bar(x):  # x in 1_O-basis
    y=toM(x); r={}
    for O,p in y.items():
        r=eadd(r,escale(bar_basis(O),pbar(p)))
    return r

def check_relations(N=6):
    # check H_s M_O formulas used for bar: 1_{(1,n)} = T_{s0} 1_{(0,n-1)}
    for n in range(1,N):
        assert T(0,{(0,n-1):ONE})=={(1,n):ONE}
        assert T(1,{(1,n-1):ONE})=={(0,n):ONE}
    # quadratic relation and omega conjugation on a test vector
    test={'P':ONE,(0,0):mono(1),(1,2):mono(-1,3),(0,3):ONE}
    for s in (0,1):
        lhs=T(s,T(s,test)); rhs=eadd(escale(T(s,test),padd(Q,ONE,-1)),escale(test,Q))
        assert lhs==rhs
    assert Tw(T(0,Tw(test)))==T(1,test)
    # bar is an involution
    for O in ['P','M',(0,0),(1,0),(0,1),(1,1),(0,2),(1,3)]:
        assert bar(bar(M(O)))==M(O), O
    print("relations ok")

def orbits_upto(N):
    L=['P','M']
    for n in range(0,N+1):
        L+= [(0,n),(1,n)]
    return L
def leq_dim_sorted(keys): return sorted(keys,key=lambda O:-d(O))

IC={}
IC['P']=M('P'); IC['M']=M('M')
IC[(0,0)]=escale({'P':ONE,'M':ONE,(0,0):ONE},mono(1))
IC[(1,0)]=Tw(IC[(0,0)])

def kl_reduce(E, top):
    """E bar-invariant in 1_O basis, top orbit coefficient M_top. Subtract IC's to reach KL form.
    Returns (IC_top, multiplicities dict O->poly)."""
    mult={}
    y=toM(E)
    assert y[top]==ONE, (top,y[top])
    changed=True
    for O in leq_dim_sorted([o for o in y if o!=top]):
        p=toM(E).get(O,{})
        nonpos={e:c for e,c in p.items() if e<=0}
        if not nonpos: continue
        c=dict(nonpos)
        for e,cc in nonpos.items():
            if e<0: c[-e]=c.get(-e,0)+cc
        mult[O]=c
        E=eadd(E,escale(IC[O],c),-1)
    return E,mult

def is_KL(E,top):
    y=toM(E)
    if y.get(top)!=ONE: return False
    for O,p in y.items():
        if O==top: continue
        if any(e<=0 for e in p): return False
    return True

def fmt(p):
    if not p: return "0"
    return " + ".join(f"{c}v^{e}" for e,c in sorted(p.items()))

if __name__=="__main__":
    check_relations()
    N=int(sys.argv[1]) if len(sys.argv)>1 else 12
    assert bar(IC[(0,0)])==IC[(0,0)]
    results=[]
    for n in range(0,N):
        # IC_{(1,n+1)} from Cs0 IC_{(0,n)}
        E=Cs(0,IC[(0,n)])
        assert bar(E)==E
        ICt,mult=kl_reduce(E,(1,n+1))
        assert is_KL(ICt,(1,n+1)) and bar(ICt)==ICt
        IC[(1,n+1)]=ICt
        IC[(0,n+1)]=Tw(ICt)
        # sanity: Cs1 IC_{(1,n)} reduces the same way to IC_{(0,n+1)}
        E2=Cs(1,IC[(1,n)])
        ICt2,mult2=kl_reduce(E2,(0,n+1))
        assert ICt2==IC[(0,n+1)]
        results.append((n,mult))
        print(f"Cs0 * IC(0,{n}) = IC(1,{n+1}) + " + ", ".join(f"[{fmt(c)}] IC{O}" for O,c in mult.items()))
    # stalks of IC(0,n)
    for n in range(0,min(N,7)):
        y=toM(IC[(0,n)])
        print(f"IC(0,{n}) in M-basis:", {str(O):fmt(p) for O,p in sorted(y.items(),key=lambda t:(d(t[0]),str(t[0])))})
