"""Build the linear system of first-order consistency equations."""
import itertools, sys, time
from collections import defaultdict, deque
from pmap import from_word, ClosedComponent
from rewrite import find_matches, glue, context_matrix
from evalmap import eval_map, sfp, rotate_vec, rot_perm
from relations import build_relations

ERRDIMS = {"a": 1, "b": 1, "E2": 4, "E25": 11, "E3": 41}
ERRORDER = ["a", "b", "E2", "E25", "E3"]


def err_offsets():
    off, o = {}, 0
    for n in ERRORDER:
        off[n] = o
        o += ERRDIMS[n]
    return off, o


class System:
    def __init__(self, t, maxk=6, maxv=None, rels=None, verbose=True):
        self.t = t
        self.one = t ** 0
        self.zero = 0 * t
        self.rels = rels or build_relations(t)
        self.maxk = maxk
        self.maxv = maxv
        self.reps = {}       # code -> rep map
        self.order = []
        self.rows = []       # list of dict col->val ; columns: ('nu', code, i) or ('err', name, i)
        self.verbose = verbose
        self._strand_coords = None

    def canon(self, X):
        code, r, stab = X.canon()
        if code not in self.reps:
            Rep = X.rotated(r)
            self.reps[code] = Rep
            self.order.append(code)
            self.pending.append(code)
            # stabiliser equations
            k = X.k
            for s in stab:
                d = (s - r) % k
                if d == 0:
                    continue
                perm = rot_perm(k, d)
                # nu(Rep) - rot_d(nu(Rep)) = 0 ; rot_d(v)[perm[i]] = v[i]
                for i in range(len(sfp(k))):
                    row = defaultdict(lambda: self.zero)
                    # component perm[i] of equation: v[perm[i]] - v[i] = 0
                    row[("nu", code, perm[i])] += self.one
                    row[("nu", code, i)] -= self.one
                    self.addrow(row)
        return code, r

    def addrow(self, row):
        row = {c: v for c, v in row.items() if v != 0}
        if row:
            self.rows.append(row)

    def nu_terms(self, X, coeff, rowvecs, sign=1):
        """add coeff * nu(X) into rowvecs (list over components of dicts)"""
        k = X.k
        loops = X.loops
        if loops:
            X = type(X)(X.nodes, X.bdry, X.alpha, 0)
            coeff = coeff * (self.t - 1) ** loops
        code, r = self.canon(X)
        # nu(X) = rotate_vec(nu(Rep), -r): component perm[i] of result = nu(Rep)[i]
        perm = rot_perm(k, (-r) % k)
        for i in range(len(sfp(k))):
            rowvecs[perm[i]][("nu", code, i)] += sign * coeff

    def err_vec(self, rel, kL):
        """returns list over unknown components: (colname, vector in L coords)"""
        name = rel["err"]
        if name is None:
            return []
        if name == "a":
            return [(("err", "a", 0), eval_map(rel["rhs"][0][1], self.t))]
        if name == "b":
            return [(("err", "b", 0), eval_map(rel["rhs"][0][1], self.t))]
        n = len(sfp(kL))
        out = []
        for i in range(n):
            v = [self.zero] * n
            v[i] = self.one
            out.append((("err", name, i), v))
        return out

    def process(self, code):
        M = self.reps[code]
        k = M.k
        nv = len(M.nodes)
        for rel in self.rels:
            L = rel["L"]
            for region, legs in find_matches(M, L):
                try:
                    rowvecs = [defaultdict(lambda: self.zero) for _ in range(len(sfp(k)))]
                    self.nu_terms(M, self.one, rowvecs, +1)
                    news = []
                    for c, R in rel["rhs"]:
                        G = glue(M, region, legs, R)
                        G.code()  # may raise
                        news.append((c, G))
                    for c, G in news:
                        self.nu_terms(G, c, rowvecs, -1)
                except ClosedComponent:
                    self.skipped += 1
                    continue
                ev = self.err_vec(rel, L.k)
                if ev:
                    C = context_matrix(M, region, legs, L.k, self.t)  # cols over sf(L.k)
                    for col, vec in ev:
                        # contribution: sum_pi vec[pi] * C[pi]
                        tot = [self.zero] * len(sfp(k))
                        for pi, x in enumerate(vec):
                            if x != 0:
                                Cp = C[pi]
                                for i in range(len(tot)):
                                    tot[i] += x * Cp[i]
                        for i in range(len(tot)):
                            if tot[i] != 0:
                                rowvecs[i][col] -= tot[i]
                for rv in rowvecs:
                    self.addrow(rv)

    def run(self, seeds, maxmaps=10 ** 6):
        self.pending = deque()
        self.skipped = 0
        for X in seeds:
            try:
                if X.loops:
                    X = type(X)(X.nodes, X.bdry, X.alpha, 0)
                self.canon(X)
            except ClosedComponent:
                pass
        n = 0
        t0 = time.time()
        while self.pending:
            code = self.pending.popleft()
            self.process(code)
            n += 1
            if self.verbose and n % 200 == 0:
                print(f"  processed {n}, reps {len(self.reps)}, rows {len(self.rows)}, {time.time()-t0:.1f}s", file=sys.stderr)
            if len(self.reps) > maxmaps:
                raise RuntimeError("too many maps")
        return self