# Hecke counts for G = Res_{E/F} GL_2 acting on hermitian 2x2 matrices, E = F_{q^2}((t)), F = F_q((t)).
# (T_lam * phi_0)(x) = #{ gK in K t^lam K / K : g^{-1} x g^{-*} in GL_2(O_E) }.
import itertools, sys
q = int(sys.argv[1]) if len(sys.argv) > 1 else 3
# non-square n0 mod q
n0 = next(a for a in range(2, q) if pow(a, (q-1)//2, q) == q-1)
# elements of F_{q^2} = F_q[i], i^2 = n0, as pairs (a,b)
def fadd(x, y): return ((x[0]+y[0]) % q, (x[1]+y[1]) % q)
def fneg(x): return ((-x[0]) % q, (-x[1]) % q)
def fmul(x, y): return ((x[0]*y[0] + n0*x[1]*y[1]) % q, (x[0]*y[1] + x[1]*y[0]) % q)
def fconj(x): return (x[0], (-x[1]) % q)
FQ2 = [(a, b) for a in range(q) for b in range(q)]
# Laurent polynomials: dict exp -> nonzero field element
def padd(P, Q):
    R = dict(P)
    for e, c in Q.items():
        s = fadd(R.get(e, (0, 0)), c)
        if s == (0, 0): R.pop(e, None)
        else: R[e] = s
    return R
def pmul(P, Q):
    R = {}
    for e1, c1 in P.items():
        for e2, c2 in Q.items():
            e = e1 + e2; s = fadd(R.get(e, (0, 0)), fmul(c1, c2))
            if s == (0, 0): R.pop(e, None)
            else: R[e] = s
    return R
def pneg(P): return {e: fneg(c) for e, c in P.items()}
def pconj(P): return {e: fconj(c) for e, c in P.items()}
def mono(e, c=(1, 0)): return {e: c}
def val(P): return min(P) if P else 10**9
def mmul(A, B):
    return [[padd(pmul(A[r][0], B[0][s]), pmul(A[r][1], B[1][s])) for s in range(2)] for r in range(2)]
def mstar(A):  # conjugate transpose
    return [[pconj(A[s][r]) for s in range(2)] for r in range(2)]
def reps(lam):
    l1, l2 = lam; tot = l1 + l2
    out = []
    for i in range(l2, tot - l2 + 1):
        j = tot - i
        exps = list(range(l2, i))
        for coeffs in itertools.product(FQ2, repeat=len(exps)):
            b = {e: c for e, c in zip(exps, coeffs) if c != (0, 0)}
            if min(i, j, val(b)) != l2: continue
            # g = [[t^i, b],[0, t^j]], g^{-1} = [[t^-i, -b t^{-i-j}],[0, t^-j]]
            ginv = [[mono(-i), pmul(pneg(b), mono(-i - j))], [{}, mono(-j)]]
            out.append(ginv)
    return out
def count(lam, c):
    x = [[mono(c[0]), {}], [{}, mono(c[1])]]
    n = 0
    R = reps(lam)
    for ginv in R:
        y = mmul(mmul(ginv, x), mstar(ginv))
        if all(val(y[r][s]) >= 0 for r in range(2) for s in range(2)):
            d = padd(pmul(y[0][0], y[1][1]), pneg(pmul(y[0][1], y[1][0])))
            if val(d) == 0: n += 1
    return n, len(R)
for lam, cs in [((1, 0), [(2, 0), (1, 1)]), ((1, 1), [(2, 2)]), ((2, 0), [(4, 0), (3, 1), (2, 2)]), ((1, -1), [(2, -2), (1, -1), (0, 0)]), ((0, 0), [(0, 0)])]:
    row = []
    for c in cs:
        n, R = count(lam, c); row.append((c, n))
    print("q=%d lam=%s #cosets=%d values:" % (q, lam, R), row)
# odd-orbit check: (1,0) has v(det) odd; every T_lam*phi_0 vanishes there for small lam
print("odd orbit (1,0):", [count(l, (1, 0))[0] for l in [(1, 0), (0, 0), (1, -1), (1, 1)]])