"""First-order errors produced by the gauge change c -> c + eps*w."""
from rewrite import context_matrix
from evalmap import eval_map, sfp, rotate_vec
from system import ERRDIMS, ERRORDER, err_offsets
from pmap import from_word


def subst_X(M, wvec, t):
    """sum over X vertices x of M of (M with x replaced by w) in coords of M"""
    k = M.k
    tot = [0 * t] * len(sfp(k))
    for vid, (tp, ds) in M.nodes.items():
        if tp != "X":
            continue
        C = context_matrix(M, {vid}, tuple(ds), 4, t)
        for pi, x in enumerate(wvec):
            if x != 0:
                for i in range(len(tot)):
                    tot[i] += x * C[pi][i]
    return tot


def gauge_errors(rels, wvec, t):
    off, n = err_offsets()
    out = [0 * t] * n
    for rel in rels:
        if rel["err"] is None:
            continue
        L = rel["L"]
        v = subst_X(L, wvec, t)
        for c, R in rel["rhs"]:
            u = subst_X(R, wvec, t)
            v = [x - c * y for x, y in zip(v, u)]
        name = rel["err"]
        if name in ("a", "b"):
            ref = eval_map(rel["rhs"][0][1], t)
            j = next(i for i, x in enumerate(ref) if x != 0)
            s = v[j] / ref[j]
            assert all(x == s * y for x, y in zip(v, ref))
            out[off[name]] = s
        else:
            for i, x in enumerate(v):
                out[off[name] + i] = x
    return out


def plus_eigen_basis(t):
    """basis of rotation-invariant 4-boxes, in sf(4) coords"""
    n = len(sfp(4))
    basis = []
    seen = set()
    for i in range(n):
        v = [0 * t] * n
        v[i] = t ** 0
        s = v
        orbit = [v]
        w = v
        for r in range(1, 4):
            w = rotate_vec(w, 4, 1)
            orbit.append(w)
        tot = [sum(xs) for xs in zip(*orbit)]
        key = tuple(int(x != 0) for x in tot)
        if key in seen:
            continue
        seen.add(key)
        basis.append(tot)
    return basis