"""Pattern matching / rewriting of planar maps, and context evaluation."""
import itertools
from pmap import PMap, from_word, smooth, ClosedComponent
from evalmap import evaluate, _merge_nodes, coords, sfp


def internal_darts(L):
    return [d for vid, (tp, ds) in L.nodes.items() for d in ds]


def find_matches(M, L):
    """Yield (region_vids, legs) for occurrences of pattern L (connected
    internal part) in host M. legs[j] = host dart corresponding to L's
    boundary position j (a dart of a matched vertex)."""
    Lb = set(L.bdry)
    Lvof = L.vof()
    Mvof = M.vof()
    vids = list(L.nodes)
    v0 = vids[0]
    r0 = L.nodes[v0][1][0]
    tp0 = L.nodes[v0][0]
    results = []
    for hv, (htp, hds) in M.nodes.items():
        if htp != tp0:
            continue
        for h in hds:
            f = {r0: h}
            vmap = {}
            stack = [r0]
            ok = True
            while stack and ok:
                d = stack.pop()
                hd = f[d]
                lv, li = Lvof[d]
                mv, mi = Mvof[hd]
                if mv == "B" or M.nodes[mv][0] != L.nodes[lv][0]:
                    ok = False
                    break
                if lv in vmap and vmap[lv] != mv:
                    ok = False
                    break
                vmap[lv] = mv
                # sigma
                for a, b in ((L.sigma(d), M.sigma(hd)),):
                    if a in f:
                        if f[a] != b:
                            ok = False
                    else:
                        f[a] = b
                        stack.append(a)
                # alpha
                ad = L.alpha[d]
                if ad in Lb:
                    continue
                bd = M.alpha[hd]
                if ad in f:
                    if f[ad] != bd:
                        ok = False
                else:
                    f[ad] = bd
                    stack.append(ad)
            if not ok:
                continue
            if len(vmap) != len(L.nodes) or len(set(vmap.values())) != len(vmap):
                continue
            # also verify internal edges map to internal edges both ways
            legs = []
            for j, bd_ in enumerate(L.bdry):
                inner = L.alpha[bd_]
                if inner in Lb:
                    ok = False  # pattern with a through strand: not supported
                    break
                legs.append(f[inner])
            if not ok:
                continue
            region = set(vmap.values())
            # sanity: legs' partners should not be internal-edge partners
            results.append((frozenset(region), tuple(legs)))
    # dedupe
    seen = set()
    out = []
    for r in results:
        if r not in seen:
            seen.add(r)
            out.append(r)
    return out


def glue(M, region, legs, R):
    """Replace the matched region by map R (same number of boundary legs)."""
    k = len(legs)
    assert R.k == k
    legset = {d: j for j, d in enumerate(legs)}
    off = M.maxdart() + 1
    Roff = R.renumbered(off)
    base = Roff.maxdart() + 1 if Roff.alpha else off + 1
    base = max(base, off + 10 + max([0] + list(Roff.alpha)) - off) if Roff.alpha else base
    base = max([off] + list(Roff.alpha) + list(Roff.bdry)) + 1
    nodes = {vid: v for vid, v in M.nodes.items() if vid not in region}
    nodes.update(Roff.nodes)
    removed = set()
    for vid in region:
        removed.update(M.nodes[vid][1])
    alpha = {d: e for d, e in M.alpha.items() if d not in removed and e not in removed}
    for d, e in Roff.alpha.items():
        alpha[d] = e
    Rb = {d: j for j, d in enumerate(Roff.bdry)}
    wire = {}
    u = [base + 2 * j for j in range(k)]
    v = [base + 2 * j + 1 for j in range(k)]
    for j in range(k):
        wire[u[j]] = v[j]
        wire[v[j]] = u[j]
    for j in range(k):
        p = M.alpha[legs[j]]
        if p in legset:
            alpha[u[j]] = u[legset[p]]
        else:
            alpha[u[j]] = p
            alpha[p] = u[j]
        r = Roff.alpha[Roff.bdry[j]]
        alpha.pop(Roff.bdry[j], None)
        if r in Rb:
            alpha[v[j]] = v[Rb[r]]
        else:
            alpha[v[j]] = r
            alpha[r] = v[j]
    newM = smooth(nodes, list(M.bdry), alpha, wire)
    newM.loops += M.loops + R.loops
    return newM


def context_matrix(M, region, legs, kL, t):
    """columns: for each sf partition pi of kL points, the coords (sf, k_M) of
    the host with the region replaced by Res(B_pi)."""
    legset = {d: j for j, d in enumerate(legs)}
    base = M.maxdart() + 1
    z = [base + j for j in range(kL)]
    removed = set()
    for vid in region:
        removed.update(M.nodes[vid][1])
    alpha0 = {d: e for d, e in M.alpha.items() if d not in removed and e not in removed}
    for j in range(kL):
        p = M.alpha[legs[j]]
        if p in legset:
            alpha0[z[j]] = z[legset[p]]
        else:
            alpha0[z[j]] = p
            alpha0[p] = z[j]
    nodes0 = _merge_nodes(M, region=region)
    cols = []
    zero = 0 * t
    for pi in sfp(kL):
        ins = [[z[j] for j in b] for b in pi]
        nodes = nodes0 + [(ds, []) for ds in ins]
        res = evaluate(nodes, alpha0, M.k, t, M.loops)
        cols.append(coords(res, M.k, zero))
    return cols