# A5: (1) fibre counts from the position formula + tree counts, as polynomials in q,
#         compared with the brute-force counts of bf_fibres.py (q = 3, 5, 7);
#     (2) degree <= f (semismallness) and leading coefficient at relevant pairs;
#     (3) branching SO5 x SO5 -> iota(SL2^3) computed from weight multiplicities
#         (Freudenthal), compared with the pair count and with the leading coefficients;
#     (4) relevance of every label, fusion check.
import json, itertools, cmath, random
from fractions import Fraction as Fr

# ---------- polynomials in q: dict power -> int ----------
def padd(a, b, s=1):
    r = dict(a)
    for k, v in b.items():
        r[k] = r.get(k, 0) + s * v
        if r[k] == 0: del r[k]
    return r
def pmul(a, b):
    r = {}
    for k1, v1 in a.items():
        for k2, v2 in b.items():
            r[k1 + k2] = r.get(k1 + k2, 0) + v1 * v2
    return {k: v for k, v in r.items() if v}
def peval(a, q): return sum(v * q**k for k, v in a.items())
def mono(k, c=1): return {k: c} if c else {}
ONE = mono(0)

def tree(rho, r1, r2):
    """# vertices at distance r1 from A and r2 from B, d(A,B)=rho, (q+1)-regular tree."""
    if (r1 + r2 - rho) % 2: return {}
    h = (r1 + r2 - rho) // 2
    j = r1 - h
    if h < 0 or j < 0 or j > rho: return {}
    if h == 0: return ONE
    if rho == 0: return padd(mono(h), mono(h - 1))           # (q+1) q^(h-1)
    if j == 0 or j == rho: return mono(h)                     # q^h
    return padd(mono(h), mono(h - 1), -1)                     # (q-1) q^(h-1)

def ext(d, DQ):
    if DQ <= d: return ONE
    s = DQ - d
    first = mono(1) if d > 0 else padd(mono(1), ONE)
    return pmul(first, mono(s - 1))

def pos_formula(d, u, w, sigma):
    m = max(u, w)
    return (d + m, max(u, 0) + max(w, 0) + sigma // 2 - m)

def n_formula(d, D, mu, amax):
    """# unimodular Q in W^perp with pos_L(C+Q) = mu, C at distance D from N (type t^d)."""
    if (D - d) % 2: return {}
    u = (D - d) // 2
    rho = min(D, d)
    tot = {}
    for DQ in range(d % 2, 2 * amax - d + 1, 2):
        w = (DQ - d) // 2
        rp = min(DQ, d)
        for sigma in range(abs(rho - rp), rho + rp + 1, 2):
            if pos_formula(d, u, w, sigma) != mu: continue
            tot = padd(tot, pmul(tree(rho, rp, sigma), ext(d, DQ)))
    return tot

def P_count(nu, lam, amax):
    d1, d2, e = nu; l1, l2 = lam
    tot = {}
    for D1 in range(d1 % 2, 2 * amax - d1 + 1, 2):
        n1 = n_formula(d1, D1, l1, amax)
        if not n1: continue
        for D2 in range(d2 % 2, 2 * amax - d2 + 1, 2):
            n2 = n_formula(d2, D2, l2, amax)
            if not n2: continue
            tot = padd(tot, pmul(tree(e, D1, D2), pmul(n1, n2)))
    return tot

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

# ---------- (1) compare with brute force ----------
for q, amax in [(3, 2), (5, 2), (7, 2), (3, 3)]:
    bf = json.load(open("bf_counts_q%d_a%d.json" % (q, amax)))
    bad_n = 0; tot_n = 0
    for key, cnt in bf["table"].items():
        d, D = map(int, key.split(","))
        for mukey, c in cnt.items():
            mu = tuple(map(int, mukey.split(",")))
            if mu[0] > amax: continue
            tot_n += 1
            if peval(n_formula(d, D, mu, amax), q) != c: bad_n += 1
    # also: formula predicts zero where brute force has no entry
    for d in range(0, 2 * amax + 1):
        for D in range(d % 2, 2 * amax - d + 1, 2):
            for mu in dominant(amax):
                c = bf["table"].get("%d,%d" % (d, D), {}).get("%d,%d" % mu, 0)
                if peval(n_formula(d, D, mu, amax), q) != c: bad_n += 1
    bad_P = 0; tot_P = 0
    labels = [(d1, d2, e) for d1 in range(2 * amax + 1) for d2 in range(2 * amax + 1)
              for e in range((d1 - d2) % 2, 4 * amax + 1, 2)]
    for l1 in dominant(amax):
        for l2 in dominant(amax):
            for nu in labels:
                c = bf["counts"].get("%s|%s" % (str(nu), str((l1, l2))), 0)
                p = peval(P_count(nu, (l1, l2), amax), q)
                tot_P += 1
                if p != c: bad_P += 1
    print("q=%d amax=%d: n-table mismatches %d (of %d entries); count mismatches %d (of %d pairs)"
          % (q, amax, bad_n, tot_n, bad_P, tot_P))

# ---------- (3) SO5 weight multiplicities by Freudenthal ----------
POS = [(1, -1), (1, 1), (1, 0), (0, 1)]
RHO = (Fr(3, 2), Fr(1, 2))
def ip(x, y): return x[0] * y[0] + x[1] * y[1]
def so5_weights(lam):
    a, b = lam
    lr = (lam[0] + RHO[0], lam[1] + RHO[1]); nlr = ip(lr, lr)
    ws = [(i, j) for i in range(-a, a + 1) for j in range(-a, a + 1) if abs(i) + abs(j) <= a + b]
    ws.sort(key=lambda m: -(3 * m[0] + m[1]))
    mult = {}
    for mu in ws:
        if mu == tuple(lam): mult[mu] = 1; continue
        mr = (mu[0] + RHO[0], mu[1] + RHO[1])
        den = nlr - ip(mr, mr)
        s = Fr(0)
        for al in POS:
            k = 1
            while True:
                nu = (mu[0] + k * al[0], mu[1] + k * al[1])
                if abs(nu[0]) > a or abs(nu[1]) > a: break
                s += mult.get(nu, 0) * ip(nu, al)
                k += 1
        m = 2 * s / den if den != 0 else 0
        assert m.denominator == 1
        if m: mult[mu] = int(m)
    return mult
def dimB2(lam):
    a, b = lam  # Weyl dimension formula for B2
    l1, l2 = a + Fr(3, 2), b + Fr(1, 2)
    return int(l1 * l2 * (l1 - l2) * (l1 + l2) / (Fr(3, 2) * Fr(1, 2) * 1 * 2))

AM = 3
W = {lam: so5_weights(lam) for lam in dominant(AM)}
for lam in dominant(AM):
    assert sum(W[lam].values()) == dimB2(lam), lam

def branching(l1, l2):
    m = {}
    for (i1, j1), c1 in W[l1].items():
        for (i2, j2), c2 in W[l2].items():
            key = (i1 + j1, i1 - j1 + i2 + j2, i2 - j2)
            m[key] = m.get(key, 0) + c1 * c2
    out = {}
    for (n1, n2, n3) in m:
        if min(n1, n2, n3) < 0: continue
        s = 0
        for eps in itertools.product((0, 1), repeat=3):
            s += (-1) ** sum(eps) * m.get((n1 + 2 * eps[0], n2 + 2 * eps[1], n3 + 2 * eps[2]), 0)
        if s: out[(n1, n2, n3)] = s
    return out, m

def Z(d, lam):
    a, b = lam
    return [D for D in range(0, 2 * a + 1) if (D - d) % 2 == 0 and abs(2 * b - d) <= D <= min(2 * b + d, 2 * a - d)]
def pairs(nu, lam):
    d1, d2, e = nu
    return sum(1 for D1 in Z(d1, lam[0]) for D2 in Z(d2, lam[1]) if abs(D1 - D2) <= e <= D1 + D2)

def two_rho(lam): return sum(4 * l[0] + 2 * l[1] for l in lam)
def delta(nu): return 3 * nu[0] + 3 * nu[1] + nu[2]

nbad_branch = 0; nbad_lead = 0; nbad_semi = 0; npairs = 0; nmult2 = 0
relevant = set()
labels = [(d1, d2, e) for d1 in range(2 * AM + 1) for d2 in range(2 * AM + 1)
          for e in range((d1 - d2) % 2, 4 * AM + 1, 2)]
for l1 in dominant(AM):
    for l2 in dominant(AM):
        br, _ = branching(l1, l2)
        for nu in labels:
            P = P_count(nu, (l1, l2), AM)
            c_pairs = pairs(nu, (l1, l2))
            c_br = br.get((nu[0], nu[2], nu[1]), 0)
            if c_pairs != c_br: nbad_branch += 1
            if not P:
                if c_pairs: nbad_lead += 1
                continue
            npairs += 1
            f2 = two_rho((l1, l2)) - delta(nu)
            deg = max(P)
            if 2 * deg > f2: nbad_semi += 1
            lead = P[deg] if 2 * deg == f2 else 0
            if lead != c_pairs: nbad_lead += 1
            if lead >= 1: relevant.add(nu)
            if lead >= 2: nmult2 += 1
        # branching entries outside the label window would indicate an error too
        for key in br:
            if key[0] > 2 * AM or key[2] > 2 * AM: pass
print("a_i<=3: nonempty (nu,lambda) pairs %d; semismall violations %d; "
      "leading coeff != pair count: %d; pair count != SO5-branching: %d; pairs with multiplicity>=2: %d"
      % (npairs, nbad_semi, nbad_lead, nbad_branch, nmult2))
# relevance of every label with d_i <= 3, e <= 6 using lambda with a_i <= 3
missing = [nu for nu in labels if nu[0] <= 3 and nu[1] <= 3 and nu[2] <= 6 and nu not in relevant]
print("labels with d_i<=3, e<=6 not relevant for any lambda with a_i<=3:", missing)
# A3-8(d) explicit lambdas
def a38d(nu):
    d1, d2, e = nu
    if e > d1 + d2:
        k = (e - d1 - d2) // 2; return ((k + d1, k), (d2, 0))
    if d1 % 2 and d2 % 2 and e == 0:
        return ((d1, 0), ((d1 + d2) // 2, abs(d1 - d2) // 2))
    s1 = None
    for s in range(d1 % 2, d1 + 1, 2):
        if (e - s) % 2 == (d2 % 2) and 0 <= e - s <= d2: s1 = s; break
    s2 = e - s1
    return (((d1 + s1) // 2, (d1 - s1) // 2), ((d2 + s2) // 2, (d2 - s2) // 2))
bad = []
for d1 in range(0, 5):
    for d2 in range(0, 5):
        for e in range((d1 - d2) % 2, 9, 2):
            nu = (d1, d2, e); lam = a38d(nu)
            am = max(lam[0][0], lam[1][0])
            P = P_count(nu, lam, am)
            f2 = two_rho(lam) - delta(nu)
            if not P or 2 * max(P) != f2 or P[max(P)] < 1: bad.append((nu, lam))
print("A3-8(d) explicit lambda fails for:", bad)
# example polynomials
for nu, lam in [((0, 0, 2), ((1, 1), (0, 0))), ((0, 0, 0), ((1, 0), (1, 0))), ((1, 1, 0), ((1, 0), (1, 0))),
                ((0, 0, 2), ((1, 0), (1, 0)))]:
    print(nu, lam, "count polynomial", dict(sorted(P_count(nu, lam, 2).items(), reverse=True)),
          " f =", (two_rho(lam) - delta(nu)) / 2, " c =", pairs(nu, lam))
# ---------- (4) fusion ----------
random.seed(1)
al, be = cmath.exp(0.37j) * 1.3, cmath.exp(1.1j) * 0.8
worst = 0
for l1 in dominant(AM):
    for l2 in dominant(AM):
        _, m = branching(l1, l2)
        v1 = sum(c * al**i * be**j * al**k for (i, j, k), c in m.items())
        v2 = sum(c * be**i * al**j * be**k for (i, j, k), c in m.items())
        worst = max(worst, abs(v1 - v2) / max(1, abs(v1)))
chi = lambda n, x: sum(x**(n - 2 * i) for i in range(n + 1))
print("fusion: max relative |chi_Res(a,b,a)-chi_Res(b,a,b)| over a_i<=3: %.2e; "
      "V_(0,2,0): %.3f" % (worst, abs(chi(2, be) - chi(2, al))))