# Brute-force check of the I-orbit classification and the Hecke action T_s 1_O on a finite ball of the
# (q+1)-regular tree.  Geodesic = oriented leaf-to-leaf path in the ball B(v0,R).
import sys, itertools
from collections import defaultdict, Counter
q=int(sys.argv[1]); R=int(sys.argv[2])
# vertices: tuples; root=(); children of root: (i,) i in 0..q ; children of non-root: append j in 0..q-1
verts=[()]
frontier=[()]
for depth in range(R):
    nf=[]
    for v in frontier:
        kids = [(i,) for i in range(q+1)] if v==() else [v+(j,) for j in range(q)]
        verts+=kids; nf+=kids
    frontier=nf
leaves=frontier
def parent(v): return v[:-1]
def nbrs(v):
    out=[] if v==() else [parent(v)]
    if len(v)<R:
        out+= [(i,) for i in range(q+1)] if v==() else [v+(j,) for j in range(q)]
    return out
def path(a,b):
    # tree path a->b
    # common prefix
    k=0
    while k<min(len(a),len(b)) and a[k]==b[k]: k+=1
    up=[a[:i] for i in range(len(a),k-1,-1)]
    down=[b[:i] for i in range(k+1,len(b)+1)]
    return up+down
def dist(x,y):
    k=0
    while k<min(len(x),len(y)) and x[k]==y[k]: k+=1
    return len(x)+len(y)-2*k
v0=(); v1=(0,)
def relpos(C,g):
    u0,u1=C; P=g
    d0=min(dist(u0,x) for x in P); d1=min(dist(u1,x) for x in P)
    if d0==0 and d1==0:
        return 'P' if P.index(u0)<P.index(u1) else 'M'
    if d0<d1: return (0,d0)
    return (1,d1)
geos=[path(a,b) for a in leaves for b in leaves if a!=b]
# orbit representatives relative to C0
C0=(v0,v1)
reps={}
for g in geos:
    O=relpos(C0,g)
    if O not in reps: reps[O]=g
print("orbits found:",sorted(reps,key=str))
# panels: s1 panel = chambers (v0,v') v' ~ v0 ; s0 panel = (v',v1), v' ~ v1
pan={1:[(v0,w) for w in nbrs(v0)], 0:[(w,v1) for w in nbrs(v1)]}
ok=True
sys.path.insert(0,'.')
from lv_module import T_on_basis
def evalpoly(p,qq): return sum(c*qq**(-e//2) for e,c in p.items()) if all(e%2==0 for e in p) else None
for s in (0,1):
    for Oprime,g in reps.items():
        if Oprime not in ('P','M') and Oprime[1]>R-3: continue
        cnt=Counter(relpos(C,g) for C in pan[s] if C!=C0)
        # predicted: coefficient of 1_{O'} in T_s 1_O  for all O
        for O in reps:
            if O not in ('P','M') and O[1]>R-3: continue
            pred=T_on_basis(s,O).get(Oprime,{})
            predval=sum(c*q**(-e//2) for e,c in pred.items())
            if cnt.get(O,0)!=predval:
                ok=False; print("MISMATCH s",s,"O'",Oprime,"O",O,cnt.get(O,0),predval)
print("Hecke action matches brute force:",ok)
# fibre count: f_n(gamma) = #{u in gamma : d(v0,u)<=n, = n mod 2} equals max(0,n+1-d(v0,gamma))
ok2=True
for g in geos:
    m=min(dist(v0,x) for x in g)
    for n in range(0,R-1):
        c=sum(1 for u in g if dist(v0,u)<=n and (dist(v0,u)-n)%2==0)
        if m<=R-2 and c!=max(0,n+1-m): ok2=False
print("fibre count of nu_n matches n+1-d(v0,gamma):",ok2)