#!/usr/bin/env python3
# duals.py -- group duals as seeds, for "On the maximality conjecture for quantum permutation groups"
# (Claude and Teo), Section 2.5 (also used in Section 3).
#
#   python3 duals.py            all the actions of Example 2.19 (Proposition 2.22), levels k = 4..7
#   python3 duals.py 8          the same, up to level 8 (20 to 40 seconds per action)
#   python3 duals.py checks     the validation cases (Klein group: Bell numbers; magic square: 14, 42)
#   python3 duals.py classc 8   the class (c) actions at N = 6 (Proposition 2.24), up to level 8
#
# Method (Remark "method" in Section 2.5).  For a Fourier family H = (H_1,...,H_r) of a group Gamma, the space
# Fix_{<S_N, dual Gamma>}(k) consists of the vectors xi = sum_pi x_pi xi_pi whose coefficients c_xi(a) vanish on all
# labelings a with w(a) != e (Proposition 2.14), where c_pi(a) = prod_b |H_{i_b}|^{1-|b|} if every block b of pi lies
# in one orbit and has labels multiplying to e, and 0 otherwise.  So
#       dim Fix = dim span{xi_pi} - rank (c_pi(a))_{a bad, pi},
# and the rank is computed exactly modulo the prime p = 1048573, by random compression (rank can only drop), so the
# output is an upper bound for the true dimension; a Catalan output therefore proves equality, since the true
# dimension is at least C_k.  Labelings giving the same row (same orbit pattern and same zero-sum subsets) are
# grouped, and a row is used when one labeling in its group has a nontrivial word.
import sys, itertools, time
import numpy as np

P = 1048573                       # prime, < 2^20, so that float64 matrix products of reduced residues are exact

# ----------------------------------------------------------------------------------------- groups
def closure(gens):
    """All elements of the group generated by gens (integer matrices), with the multiplication table."""
    gens = [np.array(g, dtype=np.int64) for g in gens]; d = gens[0].shape[0]
    e = np.eye(d, dtype=np.int64); key = lambda m: m.tobytes()
    els = [e]; idx = {key(e): 0}; frontier = [e]
    while frontier:
        new = []
        for x in frontier:
            for g in gens:
                y = x @ g
                if key(y) not in idx: idx[key(y)] = len(els); els.append(y); new.append(y)
        frontier = new
    n = len(els); mult = np.empty((n, n), dtype=np.int64)
    for i, x in enumerate(els):
        for j, y in enumerate(els): mult[i, j] = idx[key(x @ y)]
    return els, idx, mult

def perm(p, n):
    """Permutation matrix of a permutation given in cycle notation on {1..n}, e.g. '(123)(45)'."""
    s = list(range(n))
    for cyc in p.replace(' ', '').strip('()').split(')('):
        if not cyc: continue
        c = [int(ch) - 1 for ch in cyc]
        for a, b in zip(c, c[1:] + c[:1]): s[a] = b
    m = np.zeros((n, n), dtype=np.int64)
    for a in range(n): m[s[a], a] = 1
    return m

class Action:
    def __init__(self, name, gens, family):
        self.name = name
        allg = [g for H in family for g in H]
        self.els, self.idx, self.mult = closure(allg)
        key = lambda m: np.array(m, dtype=np.int64).tobytes()
        self.H = []
        for H in family:
            sub, _, _ = closure(H)
            self.H.append(sorted(self.idx[key(x)] for x in sub))
        for h in self.H:
            for a in h:
                for b in h: assert self.mult[a, b] == self.mult[b, a], 'blocks must be abelian'
        self.N = sum(len(h) for h in self.H); self.order = len(self.els)

# ----------------------------------------------------------------------------------------- partitions
def set_partitions(k):
    def rec(l):
        if l == 0: yield []; return
        for q in rec(l - 1):
            for i in range(len(q)): yield q[:i] + [q[i] + [l - 1]] + q[i + 1:]
            yield q + [[l - 1]]
    return [sorted(sorted(b) for b in q) for q in rec(k)]

def noncrossing(q):
    for A, B in itertools.permutations(q, 2):
        for a1, a2 in itertools.combinations(A, 2):
            for b1, b2 in itertools.combinations(B, 2):
                if a1 < b1 < a2 < b2: return False
    return True

def stirling2(k, j):
    S = [[0] * (j + 1) for _ in range(k + 1)]; S[0][0] = 1
    for n in range(1, k + 1):
        for m in range(1, j + 1): S[n][m] = m * S[n - 1][m] + S[n - 1][m - 1]
    return S[k][j]

# ----------------------------------------------------------------------------------------- the computation
def joint_dim(act, k, seed=1, maxrows_factor=4):
    t0 = time.time(); r = len(act.H); e = 0
    parts = set_partitions(k)
    cross = [q for q in parts if not noncrossing(q)]; ncr = len(cross)
    # the noncrossing columns vanish identically on bad labelings (xi_pi, pi noncrossing, is fixed by S_N^+),
    # so only the crossing columns are kept
    maxb = max(len(q) for q in cross) if cross else 1
    Bm = np.zeros((max(ncr, 1), maxb), dtype=np.int64)
    for i, q in enumerate(cross):
        for j, b in enumerate(q): Bm[i, j] = sum(1 << l for l in b)
    pc = np.array([bin(x).count('1') for x in range(1 << k)])
    dimspan = sum(stirling2(k, j) for j in range(1, min(k, act.N) + 1))
    if ncr == 0: return dimspan, {'rows': 0, 'dimspan': dimspan, 'rank': 0, 'secs': 0.0}
    cache = {}
    def classes(i, t):
        if (i, t) in cache: return cache[(i, t)]
        h = act.H[i]; cl = {}
        for tup in itertools.product(h, repeat=t):
            s = e
            for x in tup: s = act.mult[s, x]
            if s != e: continue
            zs = []
            for S in range(1, 1 << t):
                s = e
                for l in range(t):
                    if S >> l & 1: s = act.mult[s, tup[l]]
                if s == e: zs.append(S)
            cl.setdefault(tuple(zs), []).append(tup)
        cache[(i, t)] = list(cl.items()); return cache[(i, t)]
    gm = {}
    def globmasks(L, zs):
        key = (L, zs)
        if key not in gm: gm[key] = [sum(1 << L[l] for l in range(len(L)) if S >> l & 1) for S in zs]
        return gm[key]
    rows = []                        # per bad row group: list of (global leg mask, orbit) of the zero-sum subsets
    mult = act.mult
    for o in itertools.product(range(r), repeat=k):
        legs = [tuple(l for l in range(k) if o[l] == i) for i in range(r)]
        cls = [classes(i, len(legs[i])) if legs[i] else [((), [()])] for i in range(r)]
        for combo in itertools.product(*cls):
            bad = False
            for tups in itertools.product(*[c[1] for c in combo]):
                lab = [0] * k
                for i in range(r):
                    for l, x in zip(legs[i], tups[i]): lab[l] = x
                s = e
                for x in lab: s = mult[s, x]
                if s != e: bad = True; break
            if not bad: continue
            rows.append([(m, i) for i in range(r) if legs[i] for m in globmasks(legs[i], combo[i][0])])
    nb = len(rows); t1 = time.time()
    rng = np.random.default_rng(seed)
    if nb > maxrows_factor * (ncr + 20):   # a random subset of the rows: the rank can only drop, so this is still safe
        rows = [rows[i] for i in rng.choice(nb, maxrows_factor * (ncr + 20), replace=False)]
    invpow = np.array([[pow(pow(len(h), P - 2, P), j, P) for j in range(k + 1)] for h in act.H], dtype=np.int64)
    R_ = ncr + 20; M = np.zeros((R_, ncr))
    for c0 in range(0, len(rows), 512):
        batch = rows[c0:c0 + 512]; b = len(batch)
        V = np.zeros((b, 1 << k), dtype=np.int64); V[:, 0] = 1     # V[row, mask] = value if the mask is good, else 0
        for a, rw in enumerate(batch):
            for (m, i) in rw: V[a, m] = invpow[i, pc[m] - 1]
        vals = np.ones((b, ncr), dtype=np.int64)
        for j in range(maxb): vals = (vals * V[:, Bm[:, j]]) % P
        Rb = rng.integers(0, 1 << 12, size=(R_, b)).astype(np.float64)
        M = np.mod(M + Rb @ vals.astype(np.float64), P)
    rank = rank_mod_p(M)
    return dimspan - rank, {'rows': nb, 'dimspan': dimspan, 'rank': rank, 'secs': time.time() - t0, 'enum': t1 - t0}

def rank_mod_p(A, bs=64):
    """Rank modulo P, by blocked Gaussian elimination; all products are exact in float64 (entries < 2^20, inner
    dimensions <= bs)."""
    A = np.mod(A, P).astype(np.float64); n, m = A.shape; r = 0; c = 0
    while c < m and r < n:
        c1 = min(c + bs, m); w = c1 - c
        panel = A[r:, c:c1].copy(); F = np.zeros((n - r, w)); inv = []; q = 0
        for col in range(w):
            if q >= n - r: break
            nz = np.flatnonzero(panel[q:, col])
            if len(nz) == 0: continue
            pr = q + nz[0]
            if pr != q:
                panel[[q, pr]] = panel[[pr, q]]; F[[q, pr]] = F[[pr, q]]; A[[r + q, r + pr]] = A[[r + pr, r + q]]
            iv = pow(int(panel[q, col]), P - 2, P); inv.append(iv)
            panel[q] = np.mod(panel[q] * iv, P)
            f = panel[q + 1:, col].copy()
            panel[q + 1:] = np.mod(panel[q + 1:] - np.outer(f, panel[q]), P); F[q + 1:, q] = f
            q += 1
        T = A[r:, c1:]
        if T.shape[1]:
            for j in range(q):
                T[j] = np.mod(np.mod(T[j] - F[j, :j] @ T[:j], P) * inv[j], P)
            if q < n - r: T[q:] = np.mod(T[q:] - F[q:, :q] @ T[:q], P)
        r += q; c = c1
    return r

# ----------------------------------------------------------------------------------------- the actions
def actions():
    A = []
    p3 = lambda s: perm(s, 3); p4 = lambda s: perm(s, 4); p5 = lambda s: perm(s, 5); p7 = lambda s: perm(s, 7)
    A.append(Action('S3, three Z2', None, [[p3('(12)')], [p3('(23)')], [p3('(13)')]]))
    r, s = p4('(1234)'), p4('(24)')
    A.append(Action('D4, <s>, <r>=Z4', None, [[s], [r]]))
    A.append(Action('D4, <sr>, Klein {1,r2,s,sr2}', None, [[s @ r], [r @ r, s]]))
    A.append(Action('A4, <(123)>, <(234)>', None, [[p4('(123)')], [p4('(234)')]]))
    # Heisenberg group 3^{1+2}_+ on Z_3^2: x(a,b) = (a+1,b), y(a,b) = (a,b+a)
    pts = [(a, b) for a in range(3) for b in range(3)]; ix = {q: i for i, q in enumerate(pts)}
    def pm(f):
        m = np.zeros((9, 9), dtype=np.int64)
        for q in pts: m[ix[f(q)], ix[q]] = 1
        return m
    X = pm(lambda q: ((q[0] + 1) % 3, q[1])); Y = pm(lambda q: (q[0], (q[1] + q[0]) % 3))
    A.append(Action('3^{1+2}_+, <x>, <y>', None, [[X], [Y]]))
    def z7(f):
        m = np.zeros((7, 7), dtype=np.int64)
        for i in range(7): m[f(i) % 7, i] = 1
        return m
    t = z7(lambda i: i + 1); mm = z7(lambda i: 2 * i)                      # F_7 x| Z_3, with 2 of order 3 mod 7
    A.append(Action('F7 x| Z3, two complements', None, [[mm], [t @ mm @ np.linalg.matrix_power(t, 6)]]))
    r5, s5 = p5('(12345)'), p5('(25)(34)')
    A.append(Action('D5, <s>, <sr>, <sr2>', None, [[s5], [s5 @ r5], [s5 @ r5 @ r5]]))
    A.append(Action('A4, <(123)>, V4', None, [[p4('(123)')], [p4('(12)(34)'), p4('(13)(24)')]]))
    A.append(Action('D5, <s>, <r>=Z5', None, [[s5], [r5]]))
    # Q8 in SL_2(F_3) is awkward; use the regular permutation representation of the quaternions
    qs = ['1', 'i', 'j', 'k', '-1', '-i', '-j', '-k']
    tab = {('1', x): x for x in qs}
    def qmul(a, b):
        sa = a.startswith('-'); sb = b.startswith('-'); a0 = a.lstrip('-'); b0 = b.lstrip('-')
        base = {('1', '1'): '1', ('1', 'i'): 'i', ('1', 'j'): 'j', ('1', 'k'): 'k', ('i', '1'): 'i', ('j', '1'): 'j',
                ('k', '1'): 'k', ('i', 'i'): '-1', ('j', 'j'): '-1', ('k', 'k'): '-1', ('i', 'j'): 'k', ('j', 'k'): 'i',
                ('k', 'i'): 'j', ('j', 'i'): '-k', ('k', 'j'): '-i', ('i', 'k'): '-j'}[(a0, b0)]
        neg = sa ^ sb ^ base.startswith('-'); b1 = base.lstrip('-')
        return ('-' + b1) if neg else b1
    def qperm(g):
        m = np.zeros((8, 8), dtype=np.int64)
        for a in qs: m[qs.index(qmul(g, a)), qs.index(a)] = 1
        return m
    A.append(Action('Q8, <i>, <j>', None, [[qperm('i')], [qperm('j')]]))
    A.append(Action('D4, two Klein', None, [[r @ r, s], [r @ r, s @ r]]))
    return A

def checks():
    Xm = np.array([[0, 1], [1, 0]]); Zm = np.array([[1, 0], [0, -1]]); I = np.eye(2, dtype=np.int64)
    XZ = Xm @ Zm; kr = np.kron
    square = [kr(Xm, I), kr(I, Xm), kr(Xm, Xm), kr(I, Zm), kr(Zm, I), kr(Zm, Zm), kr(Xm, Zm), kr(Zm, Xm), kr(XZ, XZ)]
    klein = Action('Klein group, regular on 4 points', None, [[perm('(12)(34)', 4), perm('(13)(24)', 4)]])
    magic = Action('magic square 2^{1+4}_+, nine Z2', None, [[g] for g in square])
    return [(klein, (2, 3, 4), (2, 5, 15)), (magic, (4, 5), (14, 42))]

def classc():
    p5 = lambda s: perm(s, 5)
    return [Action('A5, <(123)>, <(345)>', None, [[p5('(123)')], [p5('(345)')]]),
            Action('A5, three Z2', None, [[p5('(23)(45)')], [p5('(14)(25)')], [p5('(12)(45)')]]),
            Action('A5, <(15)(23)>, Klein on {1,2,3,4}', None, [[p5('(15)(23)')], [p5('(12)(34)'), p5('(13)(24)')]]),
            Action('S5, <(1234)>, <(45)>', None, [[p5('(1234)')], [p5('(45)')]])]

CAT = {2: 2, 3: 5, 4: 14, 5: 42, 6: 132, 7: 429, 8: 1430}

def report(act, ks):
    out = []
    for k in ks:
        d, info = joint_dim(act, k); out.append(d)
        print(f'  {act.name:38s} |Gamma|={act.order:3d} N={act.N:2d}  k={k}: dim = {d:5d}'
              f'  (Catalan {CAT[k]})  [{info["rows"]} row groups, {info["secs"]:.0f} s]', flush=True)
    return out

if __name__ == '__main__':
    arg = sys.argv[1:] or []
    if arg and arg[0] == 'checks':
        for act, ks, exp in checks():
            got = report(act, ks); print('  ->', 'OK' if tuple(got) == exp else f'WRONG (expected {exp})')
    elif arg and arg[0] == 'classc':
        K = int(arg[1]) if len(arg) > 1 else 7
        for act in classc(): report(act, range(4, K + 1))
    else:
        K = int(arg[0]) if arg else 7
        for act in actions(): report(act, range(4, K + 1))