#!/usr/bin/env python3
# projlevels.py -- the projective level problem, Section 12 of "On the maximality conjecture for quantum permutation
# groups" (Claude and Teo): intermediate quantum groups PO_N <= H <= PO_N^+, with categories TL(2m) <= C(m) <= Br(2m)
# closed under tensor products, Temperley-Lieb maps between even numbers of strands, the two-point rotation rho^2 and
# the antilinear involution J (reversal + conjugation), but not necessarily under the one-point rotation rho.
#
#   python3 projlevels.py six [N...]      level 6, exact: rotation blocks of the uncappable space of Br(6), the quadric
#                                         on the rho^2-invariant block, its zero lines, closure test (default N = 4..10)
#   python3 projlevels.py gen [L] [N...]  the rho^2-category generated by TL and xi_kappa: dimensions at levels <= L
#                                         (default L = 10) against dim HL(2m) = m!, at the given N (default 5, 7)
#   python3 projlevels.py gennc [L]       the rho^2-category generated by TL alone (control: must stay TL)
#   python3 projlevels.py cert FAM m N    certificates with rho^2 blocks: FAM = PR (Brauer, PO_N <= H) or
#                                         PH (half-liberated, PU_N <= H), level 2m, at N  (uses evander.py)
#   python3 projlevels.py table           levels 8 and 10, PR and PH, N = 4..10, 12, 16  (hours: see below)
# Log: projlevelsout.txt (and evanderout/progress.txt for the certificates).
import sys, os, time, itertools
import numpy as np
import scipy.sparse as sp
from fractions import Fraction

HERE = os.path.dirname(os.path.abspath(__file__))
LOG = open(os.path.join(HERE, 'projlevelsout.txt'), 'a')
def out(s=''):
    print(s, flush=True); LOG.write(s + '\n'); LOG.flush()

# ------------------------------------------------------------------------------------------------ pairings
_P = {}
def pairings(n, fam='BR'):
    """Pairings of n points as partner tuples; fam BR (all) or HL (odd-even pairs only).  Returns (list, index dict)."""
    key = (n, fam)
    if key in _P: return _P[key]
    res = []
    def rec(part, free):
        if not free: res.append(tuple(part)); return
        a = free[0]
        for b in free[1:]:
            if fam == 'HL' and (b - a) % 2 == 0: continue
            part[a] = b; part[b] = a
            rec(part, [c for c in free if c != a and c != b])
    rec([0] * n, list(range(n)))
    res.sort()
    _P[key] = (res, {p: i for i, p in enumerate(res)})
    return _P[key]

def noncrossing(p):
    n = len(p)
    return all(not (a < c < p[a] < p[c]) for a in range(n) for c in range(n) if a < p[a] and c < p[c])

def halflib(p):
    return all((a - p[a]) % 2 for a in range(len(p)))

def rot(p, s=1):
    n = len(p); q = [0] * n
    for i in range(n): q[(i + s) % n] = (p[i] + s) % n
    return tuple(q)

def rev(p):
    n = len(p); q = [0] * n
    for i in range(n): q[n - 1 - i] = n - 1 - p[i]
    return tuple(q)

def tensor_p(p, q):
    n = len(p); return tuple(p) + tuple(x + n for x in q)

def cup_p(p, i):
    """Insert a pair at positions i, i+1 (0 <= i <= n)."""
    n = len(p); sh = lambda x: x if x < i else x + 2
    q = [0] * (n + 2)
    for a in range(n): q[sh(a)] = sh(p[a])
    q[i] = i + 1; q[i + 1] = i
    return tuple(q)

def cap_p(p, i):
    """Contract positions i and i+1 mod n.  Returns (pairing, loops)."""
    n = len(p); j = (i + 1) % n
    if p[i] == j: loops, q = 1, {a: p[a] for a in range(n) if a not in (i, j)}
    else:
        loops = 0; q = {a: p[a] for a in range(n) if a not in (i, j)}
        q[p[i]] = p[j]; q[p[j]] = p[i]
    rem = sorted(q); idx = {a: k for k, a in enumerate(rem)}
    return tuple(idx[q[a]] for a in rem), loops

def glue_p(p, q, a):
    """p o_a q : last a points of p glued to the first a points of q (rainbow: p's point n-1-i with q's point i).
    Returns (pairing on the remaining n-a + m-a points, in the order p_0..p_{n-a-1}, q_a..q_{m-1}, loops)."""
    n, m = len(p), len(q)
    parent = list(range(n + m))
    def find(x):
        while parent[x] != x: parent[x] = parent[parent[x]]; x = parent[x]
        return x
    def union(x, y):
        rx, ry = find(x), find(y)
        if rx != ry: parent[rx] = ry
    for x in range(n): union(x, p[x])
    for x in range(m): union(n + x, n + q[x])
    for i in range(a): union(n - 1 - i, n + i)
    keep = list(range(n - a)) + list(range(n + a, n + m)); idx = {x: k for k, x in enumerate(keep)}
    groups = {}
    for x in range(n + m): groups.setdefault(find(x), []).append(x)
    res = [0] * len(keep); loops = 0
    for g in groups.values():
        kp = [x for x in g if x in idx]
        if not kp: loops += 1
        else:
            assert len(kp) == 2
            res[idx[kp[0]]] = idx[kp[1]]; res[idx[kp[1]]] = idx[kp[0]]
    return tuple(res), loops

# ------------------------------------------------------------------------------------------------ operators as sparse matrices
def op_matrix(n_in, n_out, fun, N, fam_in='BR', fam_out='BR'):
    """fun(p) -> (q, loops); matrix from the basis of level n_in to level n_out (coefficient N^loops)."""
    Pi, _ = pairings(n_in, fam_in); Po, io = pairings(n_out, fam_out)
    rows, cols, vals = [], [], []
    for c, p in enumerate(Pi):
        q, l = fun(p)
        rows.append(io[q]); cols.append(c); vals.append(float(N) ** l)
    return sp.csr_matrix((vals, (rows, cols)), shape=(len(Po), len(Pi)))

# ------------------------------------------------------------------------------------------------ level 6, exact
def level_six(Ns):
    out('=== level 6, projective: uncappable space of Br(6), rho^2 blocks, quadric, closure ===')
    P6, i6 = pairings(6); P4, i4 = pairings(4); P2, i2 = pairings(2)
    tl6 = [p for p in P6 if noncrossing(p)]; cross6 = [p for p in P6 if not noncrossing(p)]
    out(f'|Br(6)| = {len(P6)}, |TL(6)| = {len(tl6)}, |HL(6)| = {sum(1 for p in P6 if halflib(p))}')
    for N in Ns:
        # exact (Fraction) uncappable space: all 6 caps (cyclic) must kill x in Br(6) (x orthogonal to TL in the
        # coordinate sense: we use the paper's convention W = C(6) minus TL(6) with caps = 0; here as in Lemma 11.3
        # the condition is cap_i x = 0 for the component ... we follow Section 11: cap_i x = 0 exactly)
        rows = []
        for i in range(6):
            M = {}
            for c, p in enumerate(P6):
                q, l = cap_p(p, i)
                M.setdefault(q, {})[c] = Fraction(N) ** l
            rows.extend(M.values())
        # but x in W is orthogonal to TL(6) for the Gram scalar product; Lemma 11.3 gives cap_i x = 0 for such x.
        # The space U = {x in Br(6): all caps kill x} is what we compute (the paper's U(6)).
        ker = kernel(rows, len(P6))
        out(f'N = {N}: dim U(6) = {len(ker)}')
        U = np.array([[float(v) for v in b] for b in ker]).T
        # rotation on U
        R = np.zeros((15, 15))
        for p in P6: R[i6[rot(p)], i6[p]] = 1
        Q, _ = np.linalg.qr(U); C = Q.T @ R @ Q
        ev = np.linalg.eigvals(C)
        out('   eigenvalues of rho on U(6): ' + ', '.join(f'{np.angle(e)/(2*np.pi)*6:+.3f}/6' for e in sorted(ev, key=lambda z: np.angle(z))) + '  (as k in e^{2 pi i k/6})')
        ev2 = np.linalg.eigvals(C @ C)
        out('   eigenvalues of rho^2 on U(6): ' + ', '.join(f'{np.angle(e)/(2*np.pi)*3:+.3f}/3' for e in sorted(ev2, key=lambda z: np.angle(z))))
        # rho-invariant block, exactly: solve (R - I) x = 0 within U
        inv = [b for b in rational_eigenspace(ker, P6, i6, N)]
        out(f'   rho-invariant block: dim {len(inv)}; rho^2-invariant block: dim {sum(1 for e in ev2 if abs(e-1)<1e-8)}')
        # quadric: crossing coefficient of x o_4 x in Br(4): the only crossing pairing of 4 points
        cr4 = [p for p in P4 if not noncrossing(p)][0]
        B = [[Fraction(0)] * 15 for _ in range(15)]
        for p in P6:
            for q in P6:
                r, l = glue_p(p, rev(q), 4)     # x o_4 Jx, J = reversal on real vectors
                if r == cr4: B[i6[p]][i6[q]] += Fraction(N) ** l
        Qf = [[sum(inv[i][r] * B[r][c] * inv[j][c] for r in range(15) for c in range(15)) for j in range(len(inv))] for i in range(len(inv))]
        out(f'   quadric q on the invariant block (basis vectors in coordinates on Br(6)):')
        for v in inv: out('      ' + str([str(x) for x in v]))
        out(f'   q = {[[str(x) for x in row] for row in Qf]}')
        if len(inv) == 2:
            a, b2, c = Qf[0][0], Qf[0][1], Qf[1][1]
            disc = b2 * b2 - a * c
            out(f'   discriminant b^2 - ac = {disc}')
            # zero lines: a s^2 + 2 b s t + c t^2 = 0
            import math
            lines = []
            if a == 0: lines = [(1, 0), (c, -2 * b2)] if b2 != 0 else [(1, 0)]
            else:
                sq = math.isqrt(int(disc.numerator)) if disc >= 0 and disc.denominator == 1 else None
                if sq is not None and sq * sq == disc.numerator:
                    for sgn in (1, -1):
                        s = (-b2 + sgn * sq) / a
                        lines.append((s, Fraction(1)))
            for (s, t) in lines:
                x = [s * inv[0][r] + t * inv[1][r] for r in range(15)]
                g = max((abs(v) for v in x if v != 0), default=1)
                supp_hl = all(halflib(P6[r]) for r in range(15) if x[r] != 0)
                # closure test: projection of x o_3 x onto Br(6) minus TL(6) (coordinates on crossing pairings) prop. to x?
                z = {}
                for p in P6:
                    for q in P6:
                        if x[i6[p]] == 0 or x[i6[q]] == 0: continue
                        r, l = glue_p(p, rev(q), 3)
                        z[r] = z.get(r, 0) + x[i6[p]] * x[i6[q]] * Fraction(N) ** l
                zc = [z.get(p, Fraction(0)) for p in P6]
                # proportional on the crossing coordinates?
                ratios = set()
                for p in cross6:
                    xi, zi = x[i6[p]], zc[i6[p]]
                    if xi == 0 and zi != 0: ratios.add('inf')
                    elif xi != 0: ratios.add(zi / xi)
                out(f'   zero line x = {[str(v) for v in x]}: supported on HL(6): {supp_hl}; '
                    f'x o_3 x crossing part proportional to x: {len(ratios) <= 1} (ratios {sorted(map(str, ratios))})')
        out()

def kernel(rows, ncols):
    M = [[Fraction(0)] * ncols for _ in rows]
    for i, r in enumerate(rows):
        for c, v in r.items(): M[i][c] = Fraction(v)
    piv = []; ri = 0
    for c in range(ncols):
        pr = next((i for i in range(ri, len(M)) if M[i][c] != 0), None)
        if pr is None: continue
        M[ri], M[pr] = M[pr], M[ri]; inv = 1 / M[ri][c]; M[ri] = [v * inv for v in M[ri]]
        for i in range(len(M)):
            if i != ri and M[i][c] != 0:
                f = M[i][c]; M[i] = [a - f * b for a, b in zip(M[i], M[ri])]
        piv.append(c); ri += 1
        if ri == len(M): break
    free = [c for c in range(ncols) if c not in piv]; basis = []
    for fc in free:
        v = [Fraction(0)] * ncols; v[fc] = Fraction(1)
        for i, pc in enumerate(piv): v[pc] = -M[i][fc]
        basis.append(v)
    return basis

def rational_eigenspace(ker, P6, i6, N):
    """Rational basis of the rho-invariant vectors inside span(ker)."""
    d = len(ker); n = len(P6)
    # x = sum c_j ker_j, (R - I) x = 0 : rows indexed by basis element
    rows = []
    for r in range(n):
        row = {}
        for j in range(d):
            # (R ker_j)[r] = ker_j[rho^{-1} index]
            val = ker[j][i6[rot(P6[r], -1)]] - ker[j][r]
            if val != 0: row[j] = val
        if row: rows.append(row)
    cs = kernel(rows, d)
    return [[sum(c[j] * ker[j][r] for j in range(d)) for r in range(n)] for c in cs]

# ------------------------------------------------------------------------------------------------ generated rho^2-category
def generated(L, N, generators, fam_ambient='BR', verbose=True):
    """Closure of the given level -> list of vectors (dicts pairing -> coef), together with TL, under: tensor products,
    cups at all positions, caps at all positions, rho^2, reversal, gluings x o_a y, within levels <= L (even levels).
    Returns {level: dim}.  Floating point at the given N, ranks by SVD."""
    levels = list(range(2, L + 1, 2))
    S = {}
    for k in levels:
        Pk, ik = pairings(k, fam_ambient)
        vecs = []
        for p in Pk:
            if noncrossing(p):
                v = np.zeros(len(Pk)); v[ik[p]] = 1; vecs.append(v)
        for g in generators.get(k, []):
            v = np.zeros(len(Pk))
            for p, c in g.items(): v[ik[p]] += c
            vecs.append(v)
        S[k] = basis(np.array(vecs).T)
    # precompute operators
    ops = {}
    def getop(key, builder):
        if key not in ops: ops[key] = builder()
        return ops[key]
    changed = True; rounds = 0
    while changed:
        changed = False; rounds += 1
        for k in levels:
            Pk, ik = pairings(k, fam_ambient); new = [S[k]]
            # rho^2 and reversal
            new.append(getop(('rot2', k), lambda: op_matrix(k, k, lambda p: (rot(p, 2), 0), N)) @ S[k])
            new.append(getop(('rev', k), lambda: op_matrix(k, k, lambda p: (rev(p), 0), N)) @ S[k])
            # cups from k-2
            if k - 2 in S:
                for i in range(k - 1):
                    new.append(getop(('cup', k - 2, i), lambda i=i: op_matrix(k - 2, k, lambda p: (cup_p(p, i), 0), N)) @ S[k - 2])
            # caps from k+2 (including cyclic)
            if k + 2 in S:
                for i in range(k + 2):
                    new.append(getop(('cap', k + 2, i), lambda i=i: op_matrix(k + 2, k, lambda p: cap_p(p, i), N)) @ S[k + 2])
            # tensor products
            for a in levels:
                b = k - a
                if b in S and a <= b:
                    new.append(tensor_span(S[a], a, S[b], b, N, fam_ambient))
            # gluings x o_g y with a + b - 2g = k, g >= 1
            for a in levels:
                for b in levels:
                    g2 = a + b - k
                    if g2 <= 0 or g2 % 2 or a < b: continue
                    g = g2 // 2
                    if g >= a or g >= b: continue
                    if len(pairings(a, fam_ambient)[0]) * len(pairings(b, fam_ambient)[0]) > 2000000: continue   # table too big
                    new.append(glue_span(S[a], a, S[b], b, g, N, fam_ambient))
            M = np.hstack(new)
            B = basis(M)
            if B.shape[1] > S[k].shape[1]:
                changed = True; S[k] = B
        if verbose: out(f'   round {rounds}: dims ' + ', '.join(f'{k}:{S[k].shape[1]}' for k in levels))
    return {k: S[k].shape[1] for k in levels}, S

def basis(M, tol=1e-9):
    if M.shape[1] == 0: return M
    if not np.all(np.isfinite(M)): raise ValueError('basis: non-finite entries (overflow?)')
    try:
        u, s, vt = np.linalg.svd(M, full_matrices=False)
    except np.linalg.LinAlgError:          # numpy's gesdd can fail to converge (seen on macOS, numpy 2.x); use gesvd
        import scipy.linalg as sla
        u, s, vt = sla.svd(M, full_matrices=False, lapack_driver='gesvd')
    r = int((s > tol * max(1.0, s[0])).sum())
    return u[:, :r]

_TT = {}
def tensor_table(a, b, fam):
    key = ('t', a, b, fam)
    if key not in _TT:
        Pa, _ = pairings(a, fam); Pb, _ = pairings(b, fam); Pk, ik = pairings(a + b, fam)
        _TT[key] = np.array([[ik[tensor_p(p, q)] for q in Pb] for p in Pa], dtype=np.int64)
    return _TT[key]

def glue_table(a, b, g, fam):
    key = ('g', a, b, g, fam)
    if key not in _TT:
        Pa, _ = pairings(a, fam); Pb, _ = pairings(b, fam); Pk, ik = pairings(a + b - 2 * g, fam)
        T = np.empty((len(Pa), len(Pb)), dtype=np.int64); L = np.empty((len(Pa), len(Pb)), dtype=np.int8)
        for i, p in enumerate(Pa):
            for j, q in enumerate(Pb):
                r, l = glue_p(p, q, g); T[i, j] = ik[r]; L[i, j] = l
        _TT[key] = (T, L)
    return _TT[key]

def bilinear_span(A, B, T, Lp, nk, N):
    """Span of all images of pairs of columns under the bilinear map e_i (x) e_j -> N^L[i,j] e_T[i,j]."""
    cols = []; acc = np.zeros((nk, 0)); ncols = 0
    nb = B.shape[0]
    W = (float(N) ** Lp.astype(np.float64))
    for x in A.T:
        nz = np.nonzero(np.abs(x) > 1e-12)[0]
        if len(nz) == 0: continue
        data = (x[nz, None] * W[nz, :]).ravel(); rows = T[nz, :].ravel(); cc = np.tile(np.arange(nb), len(nz))
        M = sp.csr_matrix((data, (rows, cc)), shape=(nk, nb))
        cols.append(M @ B); ncols += B.shape[1]
        if ncols > 1200:            # reduce the running span, to keep memory bounded
            acc = basis(np.hstack([acc] + cols)); cols = []; ncols = acc.shape[1]
    return basis(np.hstack([acc] + cols)) if cols else acc

def tensor_span(A, a, B, b, N, fam):
    Pk, _ = pairings(a + b, fam); T = tensor_table(a, b, fam)
    return bilinear_span(A, B, T, np.zeros_like(T, dtype=np.int8), len(Pk), N)

def glue_span(A, a, B, b, g, N, fam):
    Pk, _ = pairings(a + b - 2 * g, fam); T, L = glue_table(a, b, g, fam)
    return bilinear_span(A, B, T, L, len(Pk), N)

def gen_main(L, Ns, with_kappa=True):
    kappa = (3, 4, 5, 0, 1, 2)      # (14)(25)(36) in 0-based: 0-3, 1-4, 2-5
    out(f'=== rho^2-category generated by TL{" and xi_kappa" if with_kappa else ""}, levels <= {L} ===')
    import math
    for N in Ns:
        t0 = time.time()
        dims, S = generated(L, N, {6: [{kappa: 1.0}]} if with_kappa else {})
        out(f'N = {N}: ' + ', '.join(f'level {k}: {dims[k]} (HL {math.factorial(k//2)}, TL {catalan(k//2)}, Br {dfact(k-1)})' for k in dims) + f'  ({time.time()-t0:.0f}s)')
        # is the generated space inside HL?
        for k in dims:
            Pk, ik = pairings(k)
            offhl = [ik[p] for p in Pk if not halflib(p)]
            leak = np.abs(S[k][offhl, :]).max() if offhl and S[k].shape[1] else 0.0
            if leak > 1e-9: out(f'   level {k}: components outside HL (max {leak:.2e})')

def catalan(m):
    from math import comb
    return comb(2 * m, m) // (m + 1)
def dfact(n):
    r = 1
    while n > 1: r *= n; n -= 2
    return r

# ------------------------------------------------------------------------------------------------ certificates via evander
def cert(fam, m, N, hours=3.0):
    """Certificates with rho^2 blocks, through evander.py (families PR and PH)."""
    sys.path.insert(0, HERE)
    import evander
    evander.OUT = os.path.join(HERE, 'evanderout')
    return evander.run_item(fam, m, N, time.time() + hours * 3600)

if __name__ == '__main__':
    args = sys.argv[1:]
    out(f'--- projlevels.py {" ".join(args)}  {time.strftime("%d %b %H:%M")} ---')
    if not args: print(open(__file__).read().split('import sys')[0]); sys.exit()
    if args[0] == 'six': level_six([int(a) for a in args[1:]] or list(range(4, 11)))
    elif args[0] == 'gen': gen_main(int(args[1]) if len(args) > 1 else 10, [int(a) for a in args[2:]] or [5, 7])
    elif args[0] == 'gennc': gen_main(int(args[1]) if len(args) > 1 else 10, [5], with_kappa=False)
    elif args[0] == 'cert':
        r = cert(args[1], int(args[2]), int(args[3]))
        out(f'{args[1]} level {2*int(args[2])} N={args[3]}: {r}')
    elif args[0] == 'table':
        import json
        for m in (4, 5):
            for fam in ('PR', 'PH'):
                for N in (4, 5, 6, 7, 8, 9, 10, 12, 16):
                    if N < m: continue
                    r = cert(fam, m, N)
                    out(f'{fam} level {2*m} N={N}: dim U {r["dimU"] if r else "?"}, {r["status"] if r else "?"}, blocks ' +
                        (', '.join(f'k={k}: dim {v["dim"]} margin {v["margin"]:.1e}' for k, v in r['blocks'].items()) if r else ''))