"""Evaluate planar maps in S_t (P on every edge) via a frontier algorithm."""
from collections import defaultdict
from deligne import canon, sf_partitions


def _merge_nodes(M, region=None, inserts=None):
    """Return list of merge nodes: (darts, externals).  X vertices become two
    2-valent merge nodes.  region: set of vids to drop.  inserts: list of
    extra merge nodes (darts) (the darts must appear in alpha2)."""
    out = []
    for vid, (tp, ds) in M.nodes.items():
        if region and vid in region:
            continue
        if tp == "T":
            out.append((list(ds), []))
        elif tp == "X":
            out.append(([ds[0], ds[2]], []))
            out.append(([ds[1], ds[3]], []))
        else:
            raise ValueError(tp)
    for j, d in enumerate(M.bdry):
        out.append(([d], [("b", j)]))
    if inserts:
        for ds in inserts:
            out.append((list(ds), []))
    return out


def evaluate(nodes, alpha, k, t, loops=0):
    """nodes: list of (darts, externals); alpha: edge involution on all darts.
    Returns dict {cyclic partition of range(k): coeff}."""
    one = t ** 0
    mt = -one / t
    owner = {}
    for i, (ds, _) in enumerate(nodes):
        for d in ds:
            owner[d] = i
    # BFS order starting from boundary nodes
    order = []
    seen = set()
    bidx = [i for i, (_, ex) in enumerate(nodes) if ex]
    queue = list(bidx) + [i for i in range(len(nodes)) if i not in bidx]
    from collections import deque
    dq = deque()
    for s in queue:
        if s in seen:
            continue
        dq.append(s)
        seen.add(s)
        while dq:
            i = dq.popleft()
            order.append(i)
            for d in nodes[i][0]:
                j = owner[alpha[d]]
                if j not in seen:
                    seen.add(j)
                    dq.append(j)
    state = {(): one}
    done = set()
    for i in order:
        ds, ex = nodes[i]
        newblock = tuple(sorted([("d", d) for d in ds] + list(ex)))
        state = {tuple(sorted(p + (newblock,))): c for p, c in state.items()}
        done.add(i)
        for d in ds:
            q = alpha[d]
            if owner[q] not in done:
                continue
            if q in ds and q < d:
                continue  # self edge handled once
            if q in ds and q == d:
                raise ValueError
            ns = defaultdict(lambda: 0 * one)
            ld, lq = ("d", d), ("d", q)
            for p, c in state.items():
                blocks = [list(b) for b in p]
                bd = next(ii for ii, b in enumerate(blocks) if ld in b)
                bq = next(ii for ii, b in enumerate(blocks) if lq in b)
                # keep
                kb = [b for ii, b in enumerate(blocks) if ii not in (bd, bq)]
                merged = [x for x in set(blocks[bd]) | set(blocks[bq]) if x not in (ld, lq)]
                fac = one
                if merged:
                    kb.append(sorted(merged))
                else:
                    fac = fac * t
                key = tuple(sorted(tuple(b) for b in kb))
                ns[key] += c * fac
                # cut
                cb = []
                fac = mt
                for ii, b in enumerate(blocks):
                    if ii == bd or ii == bq:
                        nb = [x for x in b if x not in (ld, lq)]
                        if nb:
                            cb.append(tuple(sorted(nb)))
                        else:
                            fac = fac * t
                    else:
                        cb.append(tuple(b))
                key = tuple(sorted(cb))
                ns[key] += c * fac
            state = {p: c for p, c in ns.items() if c != 0}
    res = defaultdict(lambda: 0 * one)
    lf = (t - 1) ** loops
    for p, c in state.items():
        part = canon([[x[1] for x in b] for b in p])
        assert sum(len(b) for b in part) == k
        res[part] += c * lf
    return {p: c for p, c in res.items() if c != 0}


_SF = {}


def sfp(k):
    if k not in _SF:
        _SF[k] = sf_partitions(k)
        _SF[(k, "idx")] = {p: i for i, p in enumerate(_SF[k])}
    return _SF[k]


def sfidx(k):
    sfp(k)
    return _SF[(k, "idx")]


def coords(res, k, zero):
    return [res.get(p, zero) for p in sfp(k)]


def eval_map(M, t):
    nodes = _merge_nodes(M)
    return coords(evaluate(nodes, M.alpha, M.k, t, M.loops), M.k, 0 * t)


def rot_perm(k, r):
    """permutation of sf partition indices for rotation by r: position i -> i+r"""
    idx = sfidx(k)
    out = []
    for p in sfp(k):
        q = canon([[(i + r) % k for i in b] for b in p])
        out.append(idx[q])
    return out


def rotate_vec(v, k, r):
    perm = rot_perm(k, r)
    w = [None] * len(v)
    for i, j in enumerate(perm):
        w[j] = v[i]
    return w