"""Independent audit of the hermitian Hecke counts (A2-7, A2-6(5), A2-9).

G = Res_{F_{q^2}/F_q} GL_2 acting on nondegenerate hermitian 2x2 matrices over
E = F_{q^2}((t)) by g.x = g x g^*, K = GL_2(O_E), phi_0 = 1_{X(O)}.
(T_lam * phi_0)(x) = #{ gK in K t^lam K / K : g^{-1} x g^{-*} in GL_2(O_E) }.
Exact arithmetic in F_{q^2}[t, t^{-1}].
"""
import sys
from itertools import product


def make_field(q, eps):
    # F_{q^2} = F_q[s]/(s^2 - eps), elements (a, b) = a + b s
    def add(x, y):
        return ((x[0] + y[0]) % q, (x[1] + y[1]) % q)

    def neg(x):
        return ((-x[0]) % q, (-x[1]) % q)

    def mul(x, y):
        return ((x[0] * y[0] + eps * x[1] * y[1]) % q, (x[0] * y[1] + x[1] * y[0]) % q)

    def conj(x):
        return (x[0], (-x[1]) % q)

    def norm(x):
        return (x[0] * x[0] - eps * x[1] * x[1]) % q

    elems = [(a, b) for a in range(q) for b in range(q)]
    return add, neg, mul, conj, norm, elems


class LP:
    """Laurent polynomials over F_{q^2}: dict exponent -> nonzero element."""

    def __init__(self, F, d=None):
        self.F = F
        self.d = {k: v for k, v in (d or {}).items() if v != (0, 0)}

    def __add__(self, o):
        add = self.F[0]
        d = dict(self.d)
        for k, v in o.d.items():
            d[k] = add(d.get(k, (0, 0)), v)
        return LP(self.F, d)

    def __neg__(self):
        neg = self.F[1]
        return LP(self.F, {k: neg(v) for k, v in self.d.items()})

    def __sub__(self, o):
        return self + (-o)

    def __mul__(self, o):
        add, mul = self.F[0], self.F[2]
        d = {}
        for k1, v1 in self.d.items():
            for k2, v2 in o.d.items():
                d[k1 + k2] = add(d.get(k1 + k2, (0, 0)), mul(v1, v2))
        return LP(self.F, d)

    def conj(self):
        c = self.F[3]
        return LP(self.F, {k: c(v) for k, v in self.d.items()})

    def val(self):
        return min(self.d) if self.d else None


def mono(F, e, c=(1, 0)):
    return LP(F, {e: c})


def matmul(A, B):
    return [[A[i][0] * B[0][j] + A[i][1] * B[1][j] for j in range(2)] for i in range(2)]


def star(A):
    return [[A[j][i].conj() for j in range(2)] for i in range(2)]


def is_unimodular(Y):
    for row in Y:
        for e in row:
            v = e.val()
            if v is not None and v < 0:
                return False
    det = Y[0][0] * Y[1][1] - Y[0][1] * Y[1][0]
    return det.val() == 0


def coset_inverses(F, q, lam):
    """Inverses g^{-1} of Hermite representatives g of K t^lam K / K (GL_2(E))."""
    l1, l2 = lam
    m = l1 - l2
    elems = F[5]
    neg = F[1]
    out = []
    for i in range(m + 1):
        j = m - i
        # b = polynomial of degree < i, coefficients in F_{q^2}
        for coeffs in product(elems, repeat=i):
            b = LP(F, {k: coeffs[k] for k in range(i)})
            vb = b.val()
            e = min(i, j, vb if vb is not None else 10 ** 9)
            if e != 0:
                continue
            # g = t^{l2} [[t^i, b],[0, t^j]],  g^{-1} = t^{-l2} [[t^{-i}, -b t^{-i-j}],[0, t^{-j}]]
            ginv = [[mono(F, -i - l2), (-b) * mono(F, -i - j - l2)],
                    [LP(F), mono(F, -j - l2)]]
            out.append(ginv)
    return out


def hecke_value(F, ginvs, x):
    cnt = 0
    for gi in ginvs:
        Y = matmul(matmul(gi, x), star(gi))
        if is_unimodular(Y):
            cnt += 1
    return cnt


def diag(F, c1, c2):
    return [[mono(F, c1), LP(F)], [LP(F), mono(F, c2)]]


def run(q, eps):
    F = make_field(q, eps)
    print(f"=== q = {q}, F_q^2 = F_q[s]/(s^2-{eps})")
    lams = [(0, 0), (1, 0), (1, 1), (2, 0), (1, -1), (3, 0), (2, 1)]
    cache = {}
    for lam in lams:
        cache[lam] = coset_inverses(F, q, lam)
        print(f"  #K t^{lam} K/K = {len(cache[lam])}")
    # orbit representatives c = (c1 >= c2) in a box
    cs = [(c1, c2) for c1 in range(-3, 7) for c2 in range(-3, 7) if c1 >= c2]
    table = {}
    for lam in lams:
        row = {}
        for c in cs:
            v = hecke_value(F, cache[lam], diag(F, *c))
            if v:
                row[c] = v
        table[lam] = row
        print(f"  T_{lam} * phi_0 = " + " + ".join(f"{v}*1_{c}" for c, v in sorted(row.items(), reverse=True)))
    # support check: nonzero only on c <= 2 lam (dominance, same sum)
    for lam, row in table.items():
        L = 2 * lam[0], 2 * lam[1]
        for c in row:
            ok = (c[0] + c[1] == L[0] + L[1]) and c[0] <= L[0]
            assert ok, (lam, c)
    print("  support bound c <= 2*lam verified on the box; odd |c| never hit:",
          all((c[0] + c[1]) % 2 == 0 for row in table.values() for c in row))
    # rank on lower sets L_d = {c <= (2d,0), c2 >= 0}
    for d, lamset in [(1, [(1, 0), (0, 0)]), (2, [(2, 0), (1, 1)]), (3, [(3, 0), (2, 1)])]:
        Ld = [(2 * d - k, k) for k in range(0, d + 1)]
        vecs = []
        for lam in lamset:
            if lam[0] + lam[1] == d:
                vecs.append([table[lam].get(c, 0) for c in Ld])
        print(f"  d={d}: L_d = {Ld}; Hecke vectors on L_d (labels a with 2a in L_d): {vecs}")
    # SL_2: representatives t^{-1}[[t^i,b],[0,t^j]], i+j=2, have det 1; K t^(1,-1) K / K
    ginvs = cache[(1, -1)]
    vals = [hecke_value(F, ginvs, diag(F, n, -n)) for n in range(0, 4)]
    print(f"  SL_2: (T_alpha * phi_0)(diag(t^n,t^-n)), n=0..3: {vals}  [claim q(q+1), q+1, 1, 0 = {q*(q+1)}, {q+1}, 1, 0]")
    return table


if __name__ == "__main__":
    run(3, 2)
    run(5, 2)