# Independent check of the tree recursions used for GL2/T and SL2/T (A1-2, A2-8).
# Build a ball in the (q+1)-regular tree, fix an apartment (bi-infinite geodesic through the root),
# compute d(v) = distance to the apartment, and count neighbours / distance-2 vertices by d.
import sys
from collections import deque
from fractions import Fraction

def build(q, R):
    # vertices: tuples (path from root), children labelled 0..q at root, 0..q-1 elsewhere
    # apartment: the two root-children 0 and 1, then always child 0 along each branch
    adj = {(): []}
    order = [()]
    dq = deque([()])
    while dq:
        v = dq.popleft()
        if len(v) == R:
            continue
        nchild = q + 1 if len(v) == 0 else q
        for c in range(nchild):
            w = v + (c,)
            adj[w] = [v]
            adj[v].append(w)
            order.append(w)
            dq.append(w)
    def on_apartment(v):
        if len(v) == 0:
            return True
        return v[0] in (0, 1) and all(c == 0 for c in v[1:])
    # distance to apartment: number of trailing steps after leaving the apartment
    def dist_apt(v):
        k = 0
        for i in range(len(v) + 1):
            if on_apartment(v[:i]):
                k = i
        return len(v) - k
    d = {v: dist_apt(v) for v in adj}
    return adj, d

def check(q, R=7):
    adj, d = build(q, R)
    inner = [v for v in adj if len(v) <= R - 2]
    # adjacency: A f_e (v) = #{w ~ v : d(w)=e}
    A = {}
    for v in inner:
        cnt = {}
        for w in adj[v]:
            cnt[d[w]] = cnt.get(d[w], 0) + 1
        key = d[v]
        A.setdefault(key, set()).add(tuple(sorted(cnt.items())))
    # distance-2 operator on vertices of even distance from root (type 0)
    T = {}
    for v in inner:
        if len(v) % 2 or len(v) > R - 2:
            continue
        cnt = {}
        for w in adj[v]:
            for u in adj[w]:
                if u != v:
                    cnt[d[u]] = cnt.get(d[u], 0) + 1
        T.setdefault(d[v], set()).add(tuple(sorted(cnt.items())))
    return A, T

for q in (2, 3, 5):
    A, T = check(q, 7 if q < 5 else 5)
    print('q =', q)
    for k in sorted(A):
        print('  A: vertex with d=%d sees' % k, A[k])
    for k in sorted(T):
        print('  T: vertex with d=%d sees' % k, T[k])
# Expected: A row d=0: {0:2, 1:q-1}; d>=1: {d-1:1, d+1:q}
# T rows (column form of T f_e): d=0 -> {0:2,1:2(q-1),2:q(q-1)}? (vertex-view counts); compare with
# A2-8's T f_0 = 2f_0+2f_1+f_2 etc. via transposition below.

# Polynomial-in-q check of A2-8(c) recursions from T = A^2-(q+1) (vector view on f_d).
def Apoly(vec, q):
    out = {}
    for dd, c in vec.items():
        terms = {0: 2, 1: 1} if dd == 0 else ({0: q - 1, 2: 1} if dd == 1 else {dd - 1: q, dd + 1: 1})
        for e, t in terms.items():
            out[e] = out.get(e, 0) + c * t
    return out
def Tpoly(vec, q):
    a = Apoly(Apoly(vec, q), q)
    for dd, c in vec.items():
        a[dd] = a.get(dd, 0) - (q + 1) * c
    return {k: v for k, v in a.items() if v}
for q in (2, 3, 4, 5, 7):
    rows = [Tpoly({dd: 1}, q) for dd in range(5)]
    exp = [{0: 2, 1: 2, 2: 1}, {0: 2*(q-1), 1: q-2, 3: 1}, {0: q*(q-1), 2: q-1, 4: 1}, {1: q*q, 3: q-1, 5: 1}, {2: q*q, 4: q-1, 6: 1}]
    exp = [{k: v for k, v in e.items() if v} for e in exp]
    assert rows == exp, (q, rows, exp)
    # dimension count: span{T^n f0, n<=N} inside span(f0..f_{2N}); f1 not in span
    N = 6
    vecs = []
    cur = {0: 1}
    for n in range(N + 1):
        vecs.append(cur)
        cur = Tpoly(cur, q)
    assert all(max(v) == 2 * n and v[2 * n] == 1 for n, v in enumerate(vecs))
    # f1 in span? solve triangular system on even top indices
    def rank(rows, D):
        M = [[Fraction(r.get(j, 0)) for j in range(D)] for r in rows]
        rk = 0
        for c in range(D):
            piv = next((i for i in range(rk, len(M)) if M[i][c] != 0), None)
            if piv is None:
                continue
            M[rk], M[piv] = M[piv], M[rk]
            for i in range(len(M)):
                if i != rk and M[i][c] != 0:
                    f = M[i][c] / M[rk][c]
                    M[i] = [x - f * y for x, y in zip(M[i], M[rk])]
            rk += 1
        return rk
    r0 = rank(vecs, 2 * N + 1)
    r1 = rank(vecs + [{1: 1}], 2 * N + 1)
    print('q=%d: dim span{T^n f0, n<=%d} = %d, adding f1 gives %d (f1 in span: %s)' % (q, N, r0, r1, r0 == r1))
    # GL2 side: T1^2 phi0 and the 'support exactly' slip of A1-2 step 3
    t = Apoly(Apoly({0: 1}, q), q)
    t[0] -= (q + 3)
    print('   (T1^2-(q+3)) phi0 =', {k: v for k, v in t.items() if v}, ' (support omits 0)')
print('all polynomial recursions of A2-8(c) confirmed')
