"""Exact finite-field exploration of the p=3, N=6 resonance component.

The representation is W_{6;0,q} + X^+_{6;0,q} + X^+_{6;2,1}.
This script is an exploratory calculation, not a characteristic-zero rank
certificate.  The link-state construction and all defining relations are
checked before an image dimension is reported.
"""
from itertools import combinations
import argparse
import time
import numpy as np


def mm(a, b, p):
    return a @ b % p


def cup_matching(n, opens):
    opens = set(opens)
    height = 0
    heights = []
    for i in range(n):
        heights.append(height)
        height += 1 if i in opens else -1
    start = min(range(n), key=lambda i: heights[i])
    stack, pairs = [], []
    for j in range(n):
        i = (start + j) % n
        if i in opens:
            stack.append(i)
        else:
            pairs.append((stack.pop(), i))
    if not (not stack):
        raise AssertionError('Check failed: not stack')
    return pairs


def standard_zero(n, beta, alpha, p):
    states = list(combinations(range(n), n // 2))
    index = {v: i for i, v in enumerate(states)}
    dim = len(states)
    omega = np.zeros((dim, dim), dtype=np.int64)
    e = omega.copy()
    for k, opens in enumerate(states):
        rot = tuple(sorted((i - 1) % n for i in opens))
        omega[index[rot], k] = 1
        pairs = cup_matching(n, opens)
        partner = {}
        directed_winding = {}
        for a, b in pairs:
            partner[a] = b
            partner[b] = a
            directed_winding[a, b] = int(a > b)
            directed_winding[b, a] = -int(a > b)
        i, j = 0, 1
        if partner[i] == j:
            if (i, j) in pairs:
                target, coeff = opens, beta
            else:
                new_pairs = [v for v in pairs if v != (j, i)] + [(i, j)]
                target = tuple(sorted(a for a, b in new_pairs))
                coeff = alpha
        else:
            a, b = partner[i], partner[j]
            winding = directed_winding[a, i] + directed_winding[j, b]
            if winding == int(a > b):
                joined = (a, b)
            elif winding == -int(b > a):
                joined = (b, a)
            else:
                raise AssertionError((opens, pairs, a, b, winding))
            new_pairs = [v for v in pairs if i not in v and j not in v]
            new_pairs += [(i, j), joined]
            target = tuple(sorted(a for a, b in new_pairs))
            coeff = 1
        e[index[target], k] = coeff if p is None else coeff % p
    return e, omega


def xxz(n, d, q, z, p):
    plus = (n + d) // 2
    states = []
    for c in combinations(range(n), plus):
        c = set(c)
        states.append(tuple(1 if i in c else -1 for i in range(n)))
    index = {v: i for i, v in enumerate(states)}
    dim = len(states)
    omega = np.zeros((dim, dim), dtype=np.int64)
    e = omega.copy()
    qi = pow(int(q), -1, p)
    zi = pow(int(z), -1, p)
    for j, v in enumerate(states):
        omega[index[v[1:] + v[:1]], j] = zi if v[0] == 1 else z
        if v[0] != v[1]:
            switched = (v[1], v[0]) + v[2:]
            e[index[switched], j] = 1
            e[j, j] = (-q if v[0] == 1 else -qi) % p
    return e, omega


def matpower(a, n, p):
    r = np.eye(a.shape[0], dtype=np.int64)
    for _ in range(n):
        r = mm(r, a, p)
    return r


def check_relations(e, omega, n, beta, p):
    one = np.eye(e.shape[0], dtype=np.int64)
    if not (np.array_equal(matpower(omega, n, p), one)):
        raise AssertionError('Check failed: np.array_equal(matpower(omega, n, p), one)')
    oi = matpower(omega, n - 1, p)
    es = [e]
    for _ in range(n - 1):
        es.append(mm(mm(oi, es[-1], p), omega, p))
    for i in range(n):
        if not (np.array_equal(mm(es[i], es[i], p), beta * es[i] % p)):
            raise AssertionError('Check failed: np.array_equal(mm(es[i], es[i], p), beta * es[i] % p)')
        for j in range(n):
            if (i - j) % n in (1, n - 1):
                if not (np.array_equal(mm(mm(es[i], es[j], p), es[i], p), es[i])):
                    raise AssertionError('Check failed: np.array_equal(mm(mm(es[i], es[j], p), es[i], p), es[i])')
            else:
                if not (np.array_equal(mm(es[i], es[j], p), mm(es[j], es[i], p))):
                    raise AssertionError('Check failed: np.array_equal(mm(es[i], es[j], p), mm(es[j], es[i], p))')
    for rot in (omega, oi):
        product = mm(rot, e, p)
        if not (np.array_equal(matpower(product, n - 1, p), product)):
            raise AssertionError('Check failed: np.array_equal(matpower(product, n - 1, p), product)')


class Echelon:
    def __init__(self, p):
        self.p = p
        self.rows = []
        self.pivots = []

    def add(self, vector):
        v = vector.copy() % self.p
        for pivot, row in zip(self.pivots, self.rows):
            if v[pivot]:
                v = (v - int(v[pivot]) * row) % self.p
        nz = np.flatnonzero(v)
        if not len(nz):
            return None
        pivot = int(nz[0])
        v = v * pow(int(v[pivot]), -1, self.p) % self.p
        pos = np.searchsorted(self.pivots, pivot)
        self.pivots.insert(pos, pivot)
        self.rows.insert(pos, v)
        return v


def flatten(blocks):
    return np.concatenate([b.reshape(-1) for b in blocks])


def unflatten(vector, dims):
    off = 0
    out = []
    for dim in dims:
        out.append(vector[off:off + dim * dim].reshape(dim, dim))
        off += dim * dim
    return out


def image_basis(generators, p):
    dims = [b.shape[0] for b in generators[0]]
    span = Echelon(p)
    unit = flatten([np.eye(d, dtype=np.int64) for d in dims])
    basis = [span.add(unit)]
    i = 0
    while i < len(basis):
        blocks = unflatten(basis[i], dims)
        for gen in generators:
            v = flatten([mm(b, g, p) for b, g in zip(blocks, gen)])
            new = span.add(v)
            if new is not None:
                basis.append(new)
        i += 1
    return np.array(basis), span, dims


def rank(matrix, p):
    e = Echelon(p)
    for r in matrix:
        e.add(r)
    return len(e.rows)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--prime', type=int, default=7)
    args = parser.parse_args()
    p = args.prime
    roots = [q for q in range(2, p) if (q*q-q+1) % p == 0]
    if not (roots):
        raise AssertionError('Need a prime congruent to 1 modulo 6.')
    q = roots[0]
    pairs = [standard_zero(6, -1, 1, p), xxz(6, 0, q, q, p), xxz(6, 2, q, 1, p)]
    for e, om in pairs:
        check_relations(e, om, 6, -1, p)
    print({'prime': p, 'q': q, 'relations': 'verified'}, flush=True)
    start = time.time()
    basis, echelon, dims = image_basis([[v[0] for v in pairs], [v[1] for v in pairs]], p)
    print({'image_dimension': len(basis), 'seconds': time.time()-start}, flush=True)
    transposes = np.array([flatten([b.T for b in unflatten(v, dims)]) for v in basis])
    gram = basis @ transposes.T % p
    trank = rank(gram, p)
    print({'trace_form_rank': trank, 'radical_dimension': len(basis)-trank}, flush=True)
    np.savez(f'/tmp/c3_image_mod_{p}.npz', basis=basis, dims=np.array(dims), gram=gram,
             pivots=np.array(echelon.pivots), echelon=np.array(echelon.rows))


if __name__ == '__main__':
    main()