# A5: finite-field analogue of A3-4.  G = Sp4(F_q) on V = P + P^perp, H1 = SL(P) x SL(P^perp).
# Computes dim Hom_{H1}( Ind_B^G chi , 1 (x) sigma_c ) by characters (fixed flags), and compares with
# the orbit prediction  [c in {(chi1/chi2)^{+-1}}] + [c in {(chi1 chi2)^{+-1}}]
#                      + [chi1=1]([chi2=c]+[chi2=c^-1]) + [chi2=1]([chi1=c]+[chi1=c^-1]).
import sys, cmath, itertools
q = int(sys.argv[1]) if len(sys.argv) > 1 else 11
J = [[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]]
def om(x, y): return sum(x[i] * J[i][j] * y[j] for i in range(4) for j in range(4)) % q
def matvec(M, v): return [sum(M[i][j] * v[j] for j in range(len(v))) % q for i in range(len(M))]
def inv(a): return pow(a, q - 2, q)
g = next(x for x in range(2, q) if all(pow(x, (q - 1) // p, q) != 1 for p in range(2, q) if (q - 1) % p == 0 and all(p % r for r in range(2, p))))
LOG = {pow(g, k, q): k for k in range(q - 1)}
zeta = cmath.exp(2j * cmath.pi / (q - 1))
def char(a, x): return zeta ** (a * LOG[x % q])

def kernel(M):  # basis of kernel of M (rows x 4) mod q
    M = [row[:] for row in M]; n = 4; piv = []; r = 0
    for c in range(n):
        p = next((i for i in range(r, len(M)) if M[i][c] % q), None)
        if p is None: continue
        M[r], M[p] = M[p], M[r]
        iv = inv(M[r][c]); M[r] = [x * iv % q for x in M[r]]
        for i in range(len(M)):
            if i != r and M[i][c] % q:
                f = M[i][c]; M[i] = [(x - f * y) % q for x, y in zip(M[i], M[r])]
        piv.append(c); r += 1
    free = [c for c in range(n) if c not in piv]
    basis = []
    for fc in free:
        v = [0] * n; v[fc] = 1
        for i, pc in enumerate(piv): v[pc] = (-M[i][fc]) % q
        basis.append(v)
    return basis

def solve_coords(B, x):  # B: list of 4 basis vectors (columns), returns coords of x
    A = [[B[j][i] for j in range(4)] + [x[i]] for i in range(4)]
    for c in range(4):
        p = next(i for i in range(c, 4) if A[i][c] % q)
        A[c], A[p] = A[p], A[c]
        iv = inv(A[c][c]); A[c] = [y * iv % q for y in A[c]]
        for i in range(4):
            if i != c and A[i][c]:
                f = A[i][c]; A[i] = [(y - f * z) % q for y, z in zip(A[i], A[c])]
    return [A[i][4] for i in range(4)]

def lines_in(basis):
    k = len(basis); out = []
    for coeffs in itertools.product(range(q), repeat=k):
        nz = next((c for c in coeffs if c), None)
        if nz != 1: continue
        out.append([sum(c * b[i] for c, b in zip(coeffs, basis)) % q for i in range(4)])
    return out

def fixed_flags(h):
    """list of (alpha, beta) over flags (l c L) fixed by h, alpha on l, beta on L/l."""
    res = []
    for al in range(1, q):
        E = kernel([[(h[i][j] - (al if i == j else 0)) % q for j in range(4)] for i in range(4)])
        if not E: continue
        for v in lines_in(E):
            perp = kernel([[sum(v[i] * J[i][j] for i in range(4)) % q for j in range(4)]])  # 3-dim
            # basis (v, w1, w2) of perp, completed by z
            cand = [v]
            for w in perp:
                M = cand + [w]
                if len(kernel_rows(M)) == 0: cand.append(w)
                if len(cand) == 3: break
            for z in ([1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]):
                if len(kernel_rows(cand + [z])) == 0: B = cand + [z]; break
            A = []
            for wj in (B[1], B[2]):
                co = solve_coords(B, matvec(h, wj))
                assert co[3] == 0
                A.append((co[1], co[2]))
            a11, a21 = A[0]; a12, a22 = A[1]   # column j = image of w_j
            # eigenlines of [[a11,a12],[a21,a22]]
            for be in range(1, q):
                K = kernel2([[(a11 - be) % q, a12], [a21, (a22 - be) % q]])
                if K == 2: res += [(al, be)] * (q + 1)
                elif K == 1: res.append((al, be))
    return res

def kernel_rows(vecs):  # kernel of the map coefficients -> sum c_i vecs_i (linear dependence)
    M = [[vecs[j][i] for j in range(len(vecs))] for i in range(4)]
    n = len(vecs); M = [row[:] for row in M]; r = 0; piv = []
    for c in range(n):
        p = next((i for i in range(r, 4) if M[i][c] % q), None)
        if p is None: continue
        M[r], M[p] = M[p], M[r]
        iv = inv(M[r][c]); M[r] = [x * iv % q for x in M[r]]
        for i in range(4):
            if i != r and M[i][c] % q:
                f = M[i][c]; M[i] = [(x - f * y) % q for x, y in zip(M[i], M[r])]
        piv.append(c); r += 1
    return [c for c in range(n) if c not in piv]

def kernel2(M):  # dimension of kernel of 2x2 matrix
    a, b = M[0]; c, d = M[1]
    if a % q == 0 and b % q == 0 and c % q == 0 and d % q == 0: return 2
    return 1 if (a * d - b * c) % q == 0 else 0

# conjugacy classes of SL2(F_q): (rep, size, kind, data)
sq = {x * x % q for x in range(1, q)}
eps = next(x for x in range(2, q) if x not in sq)
classes = [([[1, 0], [0, 1]], 1, 'c', 1), ([[q - 1, 0], [0, q - 1]], 1, 'c', q - 1)]
for s in (1, q - 1):
    for b in (1, eps):
        classes.append(([[s, s * b % q], [0, s]], (q * q - 1) // 2, 'u', s))
for t in range(q):
    if t in (2, q - 2): continue
    disc = (t * t - 4) % q
    rep = [[0, q - 1], [1, t]]
    if disc in sq:
        s = next(x for x in range(1, q) if (x + inv(x)) % q == t)
        classes.append((rep, q * (q + 1), 's', s))
    else:
        classes.append((rep, q * (q - 1), 'e', None))
assert sum(c[1] for c in classes) == q**3 - q

def sigma_char(cl, cexp):
    rep, size, kind, data = cl
    if kind == 'c': return (q + 1) * char(cexp, data)
    if kind == 'u': return char(cexp, data)
    if kind == 's': return char(cexp, data) + char(cexp, inv(data))
    return 0

FF = {}
for i1, c1 in enumerate(classes):
    for i2, c2 in enumerate(classes):
        h = [[0] * 4 for _ in range(4)]
        for i in range(2):
            for j in range(2):
                h[i][j] = c1[0][i][j]; h[2 + i][2 + j] = c2[0][i][j]
        FF[(i1, i2)] = fixed_flags(h)
nflags = len(FF[(0, 0)])
print("q=%d, flags fixed by identity = %d (expected %d)" % (q, nflags, (q**4 - 1) // (q - 1) * (q + 1)))

def multiplicity(a1, a2, cexp):
    tot = 0
    for (i1, i2), fl in FF.items():
        if classes[i2][2] == 'e': continue
        chiInd = sum(char(a1, al) * char(a2, be) for al, be in fl)
        tot += classes[i1][1] * classes[i2][1] * chiInd * sigma_char(classes[i2], cexp).conjugate()
    return tot / (q**3 - q)**2

def predicted(a1, a2, k):
    m = q - 1; eq = lambda x, y: (x - y) % m == 0
    r = (eq(k, a1 - a2) + eq(k, a2 - a1)) + (eq(k, a1 + a2) + eq(k, -a1 - a2))
    r += eq(a1, 0) * (eq(a2, k) + eq(a2, -k)) + eq(a2, 0) * (eq(a1, k) + eq(a1, -k))
    return r

for (a1, a2) in [(1, 2), (0, 2), (1, 3), (2, 2)]:
    row = []
    allok = True
    for k in range(q - 1):
        val = multiplicity(a1, a2, k)
        n = round(val.real)
        ok = abs(val - n) < 1e-6 and n == predicted(a1, a2, k)
        allok &= ok
        row.append(n)
    print("chi=(z^%d,z^%d): dim Hom for c=z^k, k=0..%d: %s ; matches orbit prediction: %s"
          % (a1, a2, q - 2, row, allok))
