"""Truncated check of A2-4(4) for H = GL_2 and SL_2 (unitary twist), q = 3.

Hermitian 2x2 matrices over O_E / t^N (E = F_9((t)), N = 3). For each Cartan
type c = (c1 >= c2 >= 0) with c1 < N, let S_c be the set of truncated hermitian
matrices whose elementary divisors are (t^c1, t^c2) (determined mod t^N).
Check: S_c is a single orbit of GL_2(O_E/t^N) under x -> g x g^*.
For SL_2: S'_{(2,0)} = {x in S_(2,0) : det x = t^2 mod t^3} is a single orbit of
SL_2(O_E/t^N) (this is t * {hermitian, det 1, type alpha^vee}).
"""
from itertools import product
from collections import deque

q, eps, N = 3, 2, 3
# F_9 elements encoded as ints a + 3b  (a + b s, s^2 = eps)
FE = [(a, b) for b in range(q) for a in range(q)]
idx = {e: i for i, e in enumerate(FE)}


def fmul(x, y):
    (a, b), (c, d) = FE[x], FE[y]
    return idx[((a * c + eps * b * d) % q, (a * d + b * c) % q)]


def fadd(x, y):
    (a, b), (c, d) = FE[x], FE[y]
    return idx[((a + c) % q, (b + d) % q)]


def fneg(x):
    a, b = FE[x]
    return idx[((-a) % q, (-b) % q)]


def fconj(x):
    a, b = FE[x]
    return idx[(a, (-b) % q)]


M = [[fmul(x, y) for y in range(q * q)] for x in range(q * q)]
A = [[fadd(x, y) for y in range(q * q)] for x in range(q * q)]
NEG = [fneg(x) for x in range(q * q)]
CONJ = [fconj(x) for x in range(q * q)]
ZERO, ONE = idx[(0, 0)], idx[(1, 0)]


def padd(p, r):
    return tuple(A[p[i]][r[i]] for i in range(N))


def pmul(p, r):
    out = [ZERO] * N
    for i in range(N):
        if p[i] == ZERO:
            continue
        for j in range(N - i):
            out[i + j] = A[out[i + j]][M[p[i]][r[j]]]
    return tuple(out)


def pconj(p):
    return tuple(CONJ[c] for c in p)


def pneg(p):
    return tuple(NEG[c] for c in p)


def val(p):
    for i, c in enumerate(p):
        if c != ZERO:
            return i
    return N


def act(g, x):
    # x = (x11, x12, x22) hermitian; g = ((g11,g12),(g21,g22)); returns g x g^*
    x11, x12, x22 = x
    x21 = pconj(x12)
    X = ((x11, x12), (x21, x22))
    GX = [[padd(pmul(g[i][0], X[0][j]), pmul(g[i][1], X[1][j])) for j in range(2)] for i in range(2)]
    gs = ((pconj(g[0][0]), pconj(g[1][0])), (pconj(g[0][1]), pconj(g[1][1])))
    R = [[padd(pmul(GX[i][0], gs[0][j]), pmul(GX[i][1], gs[1][j])) for j in range(2)] for i in range(2)]
    return (R[0][0], R[0][1], R[1][1])


def typ(x):
    x11, x12, x22 = x
    c2 = min(val(x11), val(x12), val(x22))
    det = padd(pmul(x11, x22), pneg(pmul(x12, pconj(x12))))
    vd = val(det)
    if c2 >= N:
        return None
    # c1 = vd - c2 when vd < N; types with c1 < N need vd - c2 < N and vd known
    if c2 == 0:
        return (vd, 0) if vd < N else None
    # c2 >= 1: x = t^c2 y; y mod t^(N - c2): type of y is (c1-c2, 0)
    y = tuple(tuple(list(p[c2:]) + [ZERO] * c2) for p in x)
    Np = N - c2
    y11, y12, y22 = y
    dety = padd(pmul(y11, y22), pneg(pmul(y12, pconj(y12))))
    vdy = val(dety)
    if vdy < Np:
        return (vdy + c2, c2)
    return None


Fq = [idx[(a, 0)] for a in range(q)]
polys_O = list(product(Fq, repeat=N))
polys_OE = list(product(range(q * q), repeat=N))
mats = []
for x11 in polys_O:
    for x22 in polys_O:
        for x12 in polys_OE:
            mats.append((x11, x12, x22))
print("hermitian matrices mod t^N:", len(mats))

bytype = {}
for x in mats:
    c = typ(x)
    if c is not None and c[0] < N:
        bytype.setdefault(c, []).append(x)

zero = tuple([ZERO] * N)
one = tuple([ONE] + [ZERO] * (N - 1))
gens_add = []
for i in range(N):
    for e in (idx[(1, 0)], idx[(0, 1)]):
        p = [ZERO] * N
        p[i] = e
        gens_add.append(tuple(p))
gens_SL = []
for a in gens_add:
    gens_SL.append(((one, a), (zero, one)))
    gens_SL.append(((one, zero), (a, one)))
# diagonal generators for GL_2: diag(u,1), u in a generating set of units
units = [tuple([idx[(0, 1)]] + [ZERO] * (N - 1))]  # s
for i in range(1, N):
    for e in (idx[(1, 0)], idx[(0, 1)]):
        p = [ONE] + [ZERO] * (N - 1)
        p[i] = e
        units.append(tuple(p))
units.append(tuple([idx[(1, 1)]] + [ZERO] * (N - 1)))  # 1 + s, generic unit
gens_GL = gens_SL + [((u, zero), (zero, one)) for u in units]


def orbit(x0, gens, domain):
    seen = {x0}
    dq = deque([x0])
    while dq:
        x = dq.popleft()
        for g in gens:
            y = act(g, x)
            if y not in seen:
                seen.add(y)
                dq.append(y)
    return seen


for c in sorted(bytype):
    S = set(bytype[c])
    x0 = next(iter(sorted(S)))
    O = orbit(x0, gens_GL, S)
    print(f"GL_2 type {c}: |S_c| = {len(S)}, |orbit| = {len(O)}, single orbit: {O == S}")

# SL_2 check on det = t^2 mod t^3 inside type (2,0)
t2 = tuple([ZERO, ZERO, ONE])
S = set(x for x in bytype[(2, 0)]
        if padd(pmul(x[0], x[2]), pneg(pmul(x[1], pconj(x[1])))) == t2)
x0 = next(iter(sorted(S)))
O = orbit(x0, gens_SL, S)
print(f"SL_2 (det = t^2) type (2,0): |S| = {len(S)}, |orbit| = {len(O)}, single orbit: {O == S}")