"""Characteristic-zero span certificate for the p=3, N=6 image algebra.

Arithmetic is exact in Q(q), q^2-q+1=0.  No floating point or modular rank
test enters the reported dimension.  Sparse Gaussian elimination is applied
to the direct sum W_{6;0,q} + X^+_{6;0,q} + X^+_{6;2,1}.  The algebra span is
closed under right multiplication by e_1 and Omega, starting from the unit.
"""
from fractions import Fraction as F
from itertools import combinations
from bisect import bisect_left
from pathlib import Path
import argparse
import json
import time
import numpy as np
from compute_c3_image import standard_zero

ZERO = (0, 0)
ONE = (1, 0)
Q = (0, 1)
QI = (1, -1)


def add(x, y):
    return (x[0] + y[0], x[1] + y[1])


def neg(x):
    return (-x[0], -x[1])


def mul(x, y):
    a, b = x
    c, d = y
    return (a*c-b*d, a*d+b*c+b*d)


def inv(x):
    a, b = x
    norm = a*a+a*b+b*b
    return (F(a+b, norm), F(-b, norm))


def xxz_exact(n, d, twist_q):
    states = []
    for c in combinations(range(n), (n+d)//2):
        c = set(c)
        states.append(tuple(1 if i in c else -1 for i in range(n)))
    idx = {v:i for i,v in enumerate(states)}
    e, omega = {}, {}
    for j,v in enumerate(states):
        if twist_q == -1:
            weight = Q if v[0] == 1 else QI
        elif twist_q:
            weight = QI if v[0] == 1 else Q
        else:
            weight = ONE
        omega[idx[v[1:]+v[:1]],j] = weight
        if v[0] != v[1]:
            e[idx[(v[1],v[0])+v[2:]],j] = ONE
            e[j,j] = neg(Q if v[0] == 1 else QI)
    return len(states), e, omega


def representations(include_inverse=False):
    ew, ow = standard_zero(6, -1, 1, None)
    ed = {(int(i),int(j)): (int(ew[i,j]),0)
          for i,j in zip(*np.nonzero(ew))}
    od = {(int(i),int(j)): ONE for i,j in zip(*np.nonzero(ow))}
    reps = [(20,ed,od), xxz_exact(6,0,True), xxz_exact(6,2,False)]
    if include_inverse:
        reps.insert(2,xxz_exact(6,0,-1))
    return reps


def sparse_mm(a,b):
    byrow = {}
    for (k,j),v in b.items():
        byrow.setdefault(k,[]).append((j,v))
    out = {}
    for (i,k),v in a.items():
        for j,w in byrow.get(k,[]):
            value = add(out.get((i,j),ZERO),mul(v,w))
            if value == ZERO:
                out.pop((i,j),None)
            else:
                out[i,j] = value
    return out


def sparse_power(a,n,dim):
    out = {(i,i):ONE for i in range(dim)}
    for _ in range(n):
        out = sparse_mm(out,a)
    return out


def check_all_relations(dim,e,o):
    unit = {(i,i):ONE for i in range(dim)}
    if not (sparse_power(o,6,dim) == unit):
        raise AssertionError('Check failed: sparse_power(o,6,dim) == unit')
    oi = sparse_power(o,5,dim)
    es = [e]
    for _ in range(5):
        es.append(sparse_mm(sparse_mm(oi,es[-1]),o))
    for i in range(6):
        if not (sparse_mm(es[i],es[i]) == {key:neg(v) for key,v in es[i].items()}):
            raise AssertionError('Check failed: sparse_mm(es[i],es[i]) == {key:neg(v) for key,v in es[i].items()}')
        for j in range(6):
            if (i-j)%6 in (1,5):
                if not (sparse_mm(sparse_mm(es[i],es[j]),es[i]) == es[i]):
                    raise AssertionError('Check failed: sparse_mm(sparse_mm(es[i],es[j]),es[i]) == es[i]')
            else:
                if not (sparse_mm(es[i],es[j]) == sparse_mm(es[j],es[i])):
                    raise AssertionError('Check failed: sparse_mm(es[i],es[j]) == sparse_mm(es[j],es[i])')
    for rot in (o,oi):
        x = sparse_mm(rot,e)
        if not (sparse_power(x,5,dim) == x):
            raise AssertionError('Check failed: sparse_power(x,5,dim) == x')


class Span:
    def __init__(self):
        self.rows = []
        self.pivots = []

    def add(self,vector):
        v = dict(vector)
        for pivot,row in zip(self.pivots,self.rows):
            coeff = v.get(pivot,ZERO)
            if coeff != ZERO:
                coeff = neg(coeff)
                for col,value in row.items():
                    new = add(v.get(col,ZERO),mul(coeff,value))
                    if new == ZERO:
                        v.pop(col,None)
                    else:
                        v[col] = new
        if not v:
            return None
        pivot = min(v)
        normalizer = inv(v[pivot])
        v = {col:mul(value,normalizer) for col,value in v.items()}
        pos = bisect_left(self.pivots,pivot)
        self.pivots.insert(pos,pivot)
        self.rows.insert(pos,v)
        return v


def right_actions(reps):
    actions = [{},{}]
    offset = 0
    for dim,e,o in reps:
        for which,gen in enumerate((e,o)):
            rows = {}
            for (k,j),value in gen.items():
                rows.setdefault(k,[]).append((j,value))
            for i in range(dim):
                for k,entries in rows.items():
                    actions[which][offset+i*dim+k] = [(offset+i*dim+j,v) for j,v in entries]
        offset += dim*dim
    return actions


def right_multiply(vector,action):
    out = {}
    for col,value in vector.items():
        for target,weight in action.get(col,[]):
            new = add(out.get(target,ZERO),mul(value,weight))
            if new == ZERO:
                out.pop(target,None)
            else:
                out[target] = new
    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--include-inverse',action='store_true',
                        help='Adjoin X^+_{6;0,q^{-1}} to the generating module.')
    parser.add_argument('--out-dir',type=Path,
                        default=Path(__file__).resolve().parents[1] / 'data',
                        help='Directory for the exact basis and summary files.')
    args = parser.parse_args()
    reps = representations(args.include_inverse)
    for dim,e,o in reps:
        check_all_relations(dim,e,o)
    print('All affine Temperley--Lieb relations verified over Q(q).',flush=True)
    start = time.time()
    actions = right_actions(reps)
    identity = {}
    offset = 0
    for dim,e,o in reps:
        identity.update({offset+i*dim+i:ONE for i in range(dim)})
        offset += dim*dim
    span = Span()
    basis = [span.add(identity)]
    i = 0
    while i < len(basis):
        for action in actions:
            new = span.add(right_multiply(basis[i],action))
            if new is not None:
                basis.append(new)
        i += 1
    output = {'coefficient_field':'Q[q]/(q^2-q+1)',
              'representation_dimensions':[d for d,e,o in reps],
              'image_dimension':len(basis),
              'generator_products_checked':2*len(basis),
              'max_basis_nonzeros':max(map(len,basis)),
              'seconds':time.time()-start}
    print(json.dumps(output,indent=2),flush=True)
    destination = args.out_dir
    destination.mkdir(parents=True, exist_ok=True)
    stem = 'c3_inverse_exact' if args.include_inverse else 'c3_exact'
    with open(destination / (stem + '_result.json'),'w') as f:
        json.dump(output,f,indent=2)
    # Fractions are serialized exactly; this is an optional reusable basis.
    serial = [{str(k):[str(v[0]),str(v[1])] for k,v in row.items()} for row in basis]
    with open(destination / (stem + '_basis.json'),'w') as f:
        json.dump(serial,f,separators=(',',':'))


if __name__ == '__main__':
    main()