# A5 independent check of the fibre counts (T_lambda * phi_0)(x_nu) for
# X = SL2^3 \ Sp4 x Sp4, by brute force over F_q((t)):
# positions of lattices computed from valuations of minors (elementary divisors),
# NOT from the position formula.  Usage: python3 -I bf_fibres.py q amax
import sys, itertools, time

q = int(sys.argv[1]); AMAX = int(sys.argv[2])
INF = 10**9

# ---- Laurent polynomials over F_q as dicts {exp: coeff} ----
def lp_mono(e, c=1):
    c %= q
    return {e: c} if c else {}
def lp_add(a, b, s=1):
    r = dict(a)
    for e, c in b.items():
        v = (r.get(e, 0) + s * c) % q
        if v: r[e] = v
        else: r.pop(e, None)
    return r
def lp_mul(a, b):
    r = {}
    for e1, c1 in a.items():
        for e2, c2 in b.items():
            e = e1 + e2
            r[e] = (r.get(e, 0) + c1 * c2) % q
    return {e: c for e, c in r.items() if c}
def val(a):
    return min(a) if a else INF

# ---- unimodular lattices in a symplectic plane (basis e,f, omega(e,f)=1) ----
# HNF basis columns b1=(t^a,0), b2=(y,t^-a), y = sum_{j=-R}^{a-1} y_j t^j
def unimodular_lattices(R):
    out = []
    for a in range(-R, R + 1):
        exps = list(range(-R, a))
        for coeffs in itertools.product(range(q), repeat=len(exps)):
            y = {e: c for e, c in zip(exps, coeffs) if c}
            out.append((a, y))
    return out

def dist_to_N(a, y, d, k=0):
    # lattice phi_k(C): coordinates (x,y)->(t^-k x, t^k y); N_d = <e, t^d f>
    # A = diag(1,t^-d) * [[t^(a-k), t^-k y],[0, t^(k-a)]]
    vy = val(y) - k if y else INF
    k1 = min(a - k, vy, -a + k - d)
    return -d - 2 * k1

def position(C, Qd, d):
    # pos_{L(d)}(C + Q), L(d) = <e1, e2, t^-d e1 + f2, f1 + t^-d e2>
    (a, y), (a2, y2) = C, Qd
    M = [[lp_mono(a), y, {}, lp_mono(-d - a2, -1)],
         [{}, lp_mono(-d - a, -1), lp_mono(a2), y2],
         [{}, {}, {}, lp_mono(-a2)],
         [{}, lp_mono(-a), {}, {}]]
    v1 = min(val(M[i][j]) for i in range(4) for j in range(4))
    v2 = INF
    for i, j in itertools.combinations(range(4), 2):
        for k, l in itertools.combinations(range(4), 2):
            m = lp_add(lp_mul(M[i][k], M[j][l]), lp_mul(M[i][l], M[j][k]), -1)
            v2 = min(v2, val(m))
    lam1 = -v1; lam2 = -v2 - lam1
    return (lam1, lam2)

def dominant(amax):
    return [(a, b) for a in range(amax + 1) for b in range(a + 1)]

t0 = time.time()
# ---- (1) verify n(d,C,mu) depends only on D = dist(C,N_d), full run with R=1 ----
if AMAX >= 1:
    Cs = unimodular_lattices(1); Qs = unimodular_lattices(1)
    bad = 0
    for d in range(0, 3):
        seen = {}
        for C in Cs:
            D = dist_to_N(C[0], C[1], d)
            cnt = {}
            for Qd in Qs:
                p = position(C, Qd, d)
                if p[0] <= 1: cnt[p] = cnt.get(p, 0) + 1
            key = tuple(sorted(cnt.items()))
            if D in seen and seen[D] != key: bad += 1
            seen[D] = key
    print("depends-only-on-D check (R=1, d<=2): violations =", bad)

# ---- (2) table n[d][D][mu] using diagonal representatives C = <t^al e, t^-al f> ----
Qs = unimodular_lattices(AMAX)
table = {}
for d in range(0, 2 * AMAX + 1):
    for D in range(d % 2, 2 * AMAX - d + 1, 2):
        al = (D - d) // 2
        C = (al, {})
        assert dist_to_N(al, {}, d) == D
        cnt = {}
        for Qd in Qs:
            p = position(C, Qd, d)
            cnt[p] = cnt.get(p, 0) + 1
        table[(d, D)] = cnt
def n_tab(d, D, mu):
    if (d, D) not in table: return 0
    return table[(d, D)].get(mu, 0)

# ---- (3) counts (T_lambda*phi0)(x_nu) by summing over all C, via (D1,D2)-histograms ----
Call = unimodular_lattices(AMAX)
HIST = {}
def hist(nu, R):
    key = (nu, R)
    if key in HIST: return HIST[key]
    d1, d2, e = nu
    k = (e - d1 + d2) // 2
    h = {}
    for (a, y) in Call:
        vy = val(y) if y else INF
        if -2 * min(a, -a, vy) > 2 * R: continue
        D1 = dist_to_N(a, y, d1)
        D2 = dist_to_N(a, y, d2, k)
        h[(D1, D2)] = h.get((D1, D2), 0) + 1
    HIST[key] = h
    return h
def count(nu, lam):
    l1, l2 = lam
    d1, d2, e = nu
    return sum(m * n_tab(d1, D1, l1) * n_tab(d2, D2, l2) for (D1, D2), m in hist(nu, l1[0]).items())

def gr_size(lam):
    a, b = lam
    Wq = (1 + 1/q)**2 * (1 + 1/q**2)
    if a == 0 and b == 0: Wl = Wq
    elif b == 0 or a == b: Wl = 1 + 1/q
    else: Wl = 1
    return round(q**(4*a + 2*b) * Wq / Wl)

def vol(nu):
    d1, d2, e = nu
    v = 1 if e == 0 else (q + 1) * q**(e - 1)
    for d in (d1, d2):
        if d >= 1: v *= q**(3*d) - q**(3*d - 2)
    return v

results = {}
ok = 0; fail = 0
for l1 in dominant(AMAX):
    for l2 in dominant(AMAX):
        lam = (l1, l2)
        mass = 0
        for d1 in range(0, 2 * l1[0] + 1):
            for d2 in range(0, 2 * l2[0] + 1):
                for e in range((d1 - d2) % 2, 2 * (l1[0] + l2[0]) + 1, 2):
                    c = count((d1, d2, e), lam)
                    if c:
                        results[((d1, d2, e), lam)] = c
                        mass += c * vol((d1, d2, e))
        target = gr_size(l1) * gr_size(l2)
        if mass == target: ok += 1
        else:
            fail += 1; print("MASS FAIL", lam, mass, target)
print("q=%d amax=%d: mass identity holds for %d lambdas, fails for %d" % (q, AMAX, ok, fail))
print("example counts:",
      {str(k): v for k, v in results.items() if k[1] in [((1,1),(0,0)), ((1,0),(1,0))]})
import json
with open("bf_counts_q%d_a%d.json" % (q, AMAX), "w") as fh:
    json.dump({"table": {"%d,%d" % k: {"%d,%d" % m: c for m, c in v.items()} for k, v in table.items()},
               "counts": {"%s|%s" % (k[0], k[1]): v for k, v in results.items()}}, fh)
print("time %.1fs" % (time.time() - t0))