"""
Faster variant of selfdual_counts.py: enumerates only t-stable Lagrangians of Q = A^*/A.
Basis of Q ordered by (k, i) ascending, so t raises the index; in RREF a t-stable span
must satisfy t(r_i) in span(r_j : j > i), checked while backtracking from the last row up.
Usage: python3 -I selfdual_counts_fast.py q
"""
import itertools, sys
from collections import Counter
sys.path.insert(0, '/private/tmp/claude-501/-Users-miltonlin-Projects-Trees/75c2c3d9-ef06-4644-b59b-d12583740b0c/scratchpad/anim/work/G3')
import selfdual_counts as base

def build_sorted(q, a_list, N):
    S = base.build(q, a_list, N)
    if S is None:
        return None
    order = sorted(range(S['dimQ']), key=lambda j: (S['basisQ'][j][1], S['basisQ'][j][0]))
    pos = {j: p for p, j in enumerate(order)}
    basisQ = [S['basisQ'][j] for j in order]
    W = [[S['W'][order[u]][order[v]] for v in range(S['dimQ'])] for u in range(S['dimQ'])]
    tmap = []
    for j in order:
        t = S['tmap'][j]
        tmap.append(pos[t] if t is not None else None)
    S2 = dict(S)
    S2.update(basisQ=basisQ, W=W, tmap=tmap, idx={b: p for p, b in enumerate(basisQ)})
    return S2

def in_span_rref(v, rows, pivots, q):
    # rows in RREF with given pivots (subset); reduce v
    w = list(v)
    for r, p in zip(rows, pivots):
        if w[p]:
            f = w[p]
            w = [(x - f * y) % q for x, y in zip(w, r)]
    return not any(w)

def t_stable_lagrangians(q, S):
    dimQ, W, tmap = S['dimQ'], S['W'], S['tmap']
    D = dimQ // 2
    def tv(v):
        out = [0] * dimQ
        for j, x in enumerate(v):
            if x and tmap[j] is not None:
                out[tmap[j]] = (out[tmap[j]] + x) % q
        return out
    for pivots in itertools.combinations(range(dimQ), D):
        pset = set(pivots)
        free = [c for c in range(dimQ) if c not in pset]
        row_free = [[c for c in free if c > p] for p in pivots]
        chosen = {}   # r -> vector, filled from r = D-1 down to 0
        def rec(r):
            if r < 0:
                yield [chosen[i] for i in range(D)]
                return
            below = [chosen[i] for i in range(r + 1, D)]
            below_piv = [pivots[i] for i in range(r + 1, D)]
            for vals in itertools.product(range(q), repeat=len(row_free[r])):
                v = [0] * dimQ
                v[pivots[r]] = 1
                for c, x in zip(row_free[r], vals):
                    v[c] = x
                # rows below are in RREF w.r.t. their pivots, but v may have entries at their
                # pivot columns? No: RREF puts zeros of row r at all other pivot columns.
                if any(base.omega(W, v, w, q) for w in below):
                    continue
                if not in_span_rref(tv(v), below, below_piv, q):
                    continue
                chosen[r] = v
                yield from rec(r - 1)
                del chosen[r]
        yield from rec(D - 1)

def count(q, a_list, N=1):
    S = build_sorted(q, a_list, N)
    if S is None:
        return Counter()
    if S['dimQ'] == 0:
        return Counter({base.lattice_type([], S, q): 1})
    C = Counter()
    for L in t_stable_lagrangians(q, S):
        C[base.lattice_type(L, S, q)] += 1
    return C

if __name__ == '__main__':
    q = int(sys.argv[1]) if len(sys.argv) > 1 else 3
    print('GL6/Sp6 (fast), q=%d' % q)
    for mu in [(2, 0, 0), (1, 1, 0), (2, 1, -1), (2, 2, -2)]:
        C = count(q, list(mu))
        print('  orbit', tuple(x for x in mu for _ in (0, 1)), dict(C))
    print('  G3 closed forms (c1..c4):')
    print('    R1 (0,0,0,0,-1,-1):', (1, (q + 1) * (q * q + 1), 0, 0))
    print('    R2 (1,0,0,-1,-1,-1):', (q * (q + 1) * (q * q + 1), q ** 3 * (q + 1) ** 2 * (q * q + 1), (q + 1) ** 2, 0))
    print('    R3 (1,1,-1,-1,-1,-1):', (q ** 3 * (q + 1) * (q * q + 1), 0, q ** 3 - q, 1))
    print('  A4-10 as written:')
    print('    R2:', ((q ** 4 - 1) // (q - 1), (q + 1) * (q * q + 1) * (q * q + q), (q + 1) ** 2, 0))
    print('    R3:', ((q + 1) * (q * q + 1), 0, 0, 1))
    print('GL4/Sp4 (fast), q=%d' % q)
    for mu in [(2, 0), (1, 1)]:
        print('  orbit', (mu[0], mu[0], mu[1], mu[1]), dict(count(q, list(mu))))