# Finite checks for A^n under GL_n at Iwahori level (agent A5).
# (a) tau = t^{e_n} sigma acts on diagonal lattices t^mu O^n by mu -> (mu_2..mu_n, mu_1+1);
#     tau^k O^n = Lambda_{m,j} with k = n m + n - j, and relative dimension d = -sum(mu) = -k.
# (b) l(t^lambda) for lambda = lambda_{m,j}: index exponent of I cap t^l I t^-l in I = j(n-j).
# (c) T_{s0} * phi0 = 1_{tau^{-1}O^n} + q 1_{tau O^n} - 1_{O^n}  (n=2,3, q=3), exact enumeration
#     on the window t^{-1}O^n / t^2 O^n, coset reps of I s0 I / I given by i_c = 1 + c t E_{n1}.
import itertools, sys

def lam(m, j, n):
    return tuple([m]*j + [m+1]*(n-j))

def tau(mu):
    return tuple(list(mu[1:]) + [mu[0]+1])

ok = True
for n in range(1, 6):
    mu = tuple([0]*n)
    for k in range(0, 3*n+1):
        m, j = divmod(k, n)            # k = n m + (n - j')  -> solve below
        # find (m,j) with n m + n - j = k, 1 <= j <= n
        found = [(mm, jj) for mm in range(-1, 5) for jj in range(1, n+1) if n*mm + n - jj == k]
        assert len(found) == 1
        mm, jj = found[0]
        if mu != lam(mm, jj, n) or -sum(mu) != -k:
            ok = False; print("chain mismatch", n, k, mu, lam(mm, jj, n))
        mu = tau(mu)
print("(a) tau^k O^n = Lambda_{m,j}, k = nm+n-j, d = -k :", ok)

ok = True
for n in range(1, 7):
    for j in range(1, n+1):
        l = lam(0, j, n)
        # entry (a,b) of t^l I t^-l is scaled by t^{l_a - l_b}; I has O above/on diagonal, tO below.
        # index exponent = sum over entries of max(0, required valuation in conj - valuation in I)
        e = 0
        for a in range(n):
            for b in range(n):
                vI = 0 if a <= b else 1
                vC = vI + l[a] - l[b]
                e += max(0, vC - vI)
        if e != j*(n-j):
            ok = False; print("length mismatch", n, j, e)
print("(b) l(t^{lambda_{m,j}}) = j(n-j):", ok)

q = 3
def val(coeffs):
    # coeffs: tuple of coefficients for t^{-1}, t^0, t^1 ; returns valuation (inf as 99)
    for i, c in enumerate(coeffs):
        if c % q:
            return i - 1
    return 99

for n in (2, 3):
    bad = 0; tot = 0
    for flat in itertools.product(range(q), repeat=3*n):
        v = [flat[3*i:3*i+3] for i in range(n)]
        vals = [val(x) for x in v]
        s = 0
        for c in range(q):
            # w_n = v_n - c t v_1 ; coefficient arrays shift by one under multiplication by t
            tv1 = (0,) + v[0][:2]
            wn = tuple((v[n-1][i] - c*tv1[i]) % q for i in range(3))
            cond = val(wn) >= 1 and vals[0] >= -1 and all(vals[i] >= 0 for i in range(1, n-1))
            s += cond
        inv = 1 if (vals[0] >= -1 and all(x >= 0 for x in vals[1:])) else 0      # tau^{-1}O^n
        inl = 1 if (all(x >= 0 for x in vals[:-1]) and vals[-1] >= 1) else 0     # tau O^n
        ino = 1 if all(x >= 0 for x in vals) else 0                               # O^n
        if s != inv + q*inl - ino:
            bad += 1
        tot += 1
    print(f"(c) n={n}, q={q}: T_s0*phi0 = 1_(tau^-1 O) + q 1_(tau O) - 1_O at {tot} points, mismatches:", bad)