"""Layered string diagrams: rendering to TikZ, evaluation in S_t, and
conversion to planar maps (for isotopy checks).

A diagram is (m, layers): m = number of bottom strands; layers is a list of
layers (bottom to top), each a list of items acting on disjoint consecutive
strands.  Items:
   ('X', i)            crossing on strands i, i+1
   ('T', i)            trivalent vertex merging strands i, i+1
   ('Ts', i)           trivalent vertex splitting strand i
   ('cup', i)          new pair of strands inserted before strand i
   ('cap', i)          strands i, i+1 capped off
   ('box', i, nin, nout, label)   coupon
   ('node', i, nin, nout)         hollow vertex (one block of a partition)
Every crossing drawn is an X item; nothing else ever crosses.
"""
import sys
import os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

ARITY = {"X": (2, 2), "T": (2, 1), "Ts": (1, 2), "cup": (0, 2), "cap": (2, 0)}


def arity(it):
    if it[0] in ARITY:
        return ARITY[it[0]]
    return it[2], it[3]


def _pav(vals, weights):
    """isotonic (non-decreasing) weighted least squares"""
    blocks = []  # [value, weight, count]
    for v, w in zip(vals, weights):
        blocks.append([v, w, 1])
        while len(blocks) > 1 and blocks[-2][0] > blocks[-1][0]:
            v2, w2, c2 = blocks.pop()
            v1, w1, c1 = blocks.pop()
            W = w1 + w2
            blocks.append([(v1 * w1 + v2 * w2) / W, W, c1 + c2])
    out = []
    for v, w, c in blocks:
        out += [v] * c
    return out


def layout(cur, items):
    """cur: x positions of current strands.  Returns (inputs_target, outputs,
    geoms) where geoms lists, for each item, (item, in_xs, out_xs)."""
    N = len(cur)
    items = sorted(items, key=lambda it: (it[1], arity(it)[0] > 0))
    blocks = []  # each: dict(kind, item, nin, nout, w, in_idx)
    k = 0
    ii = 0
    while k < N or ii < len(items):
        if ii < len(items) and items[ii][1] == k:
            it = items[ii]
            nin, nout = arity(it)
            w = max(nin, nout, 1)
            blocks.append(dict(item=it, nin=nin, nout=nout, w=w, idx=list(range(k, k + nin))))
            k += nin
            ii += 1
        elif k < N:
            blocks.append(dict(item=None, nin=1, nout=1, w=1, idx=[k]))
            k += 1
        else:
            raise ValueError("item position out of range %r" % (items[ii],))
    # rigid offsets of inputs inside a block: inputs centred in the block
    g = 0.0
    vals, wts = [], []
    for j, b in enumerate(blocks):
        span = b["w"] - 1
        b["span"] = span
        b["g"] = g
        cen = span / 2.0
        b["in_off"] = [cen - (b["nin"] - 1) / 2.0 + q for q in range(b["nin"])]
        b["out_off"] = [cen - (b["nout"] - 1) / 2.0 + q for q in range(b["nout"])]
        if b["nin"]:
            d = sum(cur[i] - o for i, o in zip(b["idx"], b["in_off"])) / b["nin"]
            wt = float(b["nin"])
        else:
            # zero-input block: sit between neighbours
            left = cur[b["idx"][0] - 1] if False else None
            kk = sum(bb["nin"] for bb in blocks[:j])
            L = cur[kk - 1] if kk - 1 >= 0 else None
            R = cur[kk] if kk < N else None
            if L is not None and R is not None:
                c = (L + R) / 2.0
            elif L is not None:
                c = L + 1 + span / 2.0
            elif R is not None:
                c = R - 1 - span / 2.0
            else:
                c = 0.0
            d = c - span / 2.0
            wt = 0.05
        vals.append(d - g)
        wts.append(wt)
        g += span + 1
    ys = _pav(vals, wts)
    ins, outs, geoms = [], [], []
    for b, y in zip(blocks, ys):
        L = y + b["g"]
        ix = [L + o for o in b["in_off"]]
        ox = [L + o for o in b["out_off"]]
        ins += ix
        outs += ox
        geoms.append((b["item"], ix, ox))
    return ins, outs, geoms


def fmt(v):
    s = "%.3f" % v
    s = s.rstrip("0").rstrip(".")
    return "0" if s in ("-0", "") else s


class Render:
    def __init__(self, xs=1.0, hact=1.0, hshift=0.55):
        self.cmds = []
        self.xs, self.hact, self.hshift = xs, hact, hshift

    def P(self, x, y):
        return "(%s,%s)" % (fmt(x * self.xs), fmt(y))

    def curve(self, x0, y0, x1, y1):
        if abs(x0 - x1) < 1e-9:
            self.cmds.append("\\draw %s -- %s;" % (self.P(x0, y0), self.P(x1, y1)))
        else:
            ym = (y0 + y1) / 2
            self.cmds.append("\\draw %s .. controls %s and %s .. %s;" % (
                self.P(x0, y0), self.P(x0, ym), self.P(x1, ym), self.P(x1, y1)))


def render(m, layers, xs=0.5, hact=0.62, hshift=0.38, top_pad=0.0, labels=None):
    R = Render(xs, hact, hshift)
    cur = [float(i) for i in range(m)]
    y = 0.0
    if m == 0 and layers:
        cur = []
    if not layers:
        for x in cur:
            R.curve(x, 0.0, x, hact)
        return R.cmds, hact
    for layer in layers:
        ins, outs, geoms = layout(cur, layer)
        # pre-shift
        if len(ins) != len(cur):
            raise ValueError("inconsistent")
        if any(abs(a - b) > 1e-6 for a, b in zip(ins, cur)):
            for a, b in zip(cur, ins):
                R.curve(a, y, b, y + hshift)
            y += hshift
        h = hact
        for it, ix, ox in geoms:
            if it is None:
                R.curve(ix[0], y, ox[0], y + h)
                continue
            kind = it[0]
            if kind == "X":
                R.curve(ix[0], y, ox[1], y + h)
                R.curve(ix[1], y, ox[0], y + h)
            elif kind == "T":
                vx, vy = ox[0], y + 0.62 * h
                for x in ix:
                    R.cmds.append("\\draw %s .. controls %s .. %s;" % (R.P(x, y), R.P(x, y + 0.38 * h), R.P(vx, vy)))
                R.cmds.append("\\draw %s -- %s;" % (R.P(vx, vy), R.P(vx, y + h)))
                R.cmds.append("\\fill %s circle (1.5pt);" % R.P(vx, vy))
            elif kind == "Ts":
                vx, vy = ix[0], y + 0.38 * h
                R.cmds.append("\\draw %s -- %s;" % (R.P(vx, y), R.P(vx, vy)))
                for x in ox:
                    R.cmds.append("\\draw %s .. controls %s .. %s;" % (R.P(vx, vy), R.P(x, y + 0.62 * h), R.P(x, y + h)))
                R.cmds.append("\\fill %s circle (1.5pt);" % R.P(vx, vy))
            elif kind == "cup":
                a, b = ox
                R.cmds.append("\\draw %s .. controls %s and %s .. %s;" % (
                    R.P(a, y + h), R.P(a, y + 0.2 * h), R.P(b, y + 0.2 * h), R.P(b, y + h)))
            elif kind == "cap":
                a, b = ix
                R.cmds.append("\\draw %s .. controls %s and %s .. %s;" % (
                    R.P(a, y), R.P(a, y + 0.8 * h), R.P(b, y + 0.8 * h), R.P(b, y)))
            elif kind in ("box", "node"):
                allx = ix + ox
                if kind == "box":
                    lo, hi = min(allx) - 0.38, max(allx) + 0.38
                    if hi - lo < 1.1:
                        c = (lo + hi) / 2
                        lo, hi = c - 0.55, c + 0.55
                    y0, y1 = y + 0.2 * h, y + 0.8 * h
                    for x in ix:
                        R.cmds.append("\\draw %s -- %s;" % (R.P(x, y), R.P(x, y0)))
                    for x in ox:
                        R.cmds.append("\\draw %s -- %s;" % (R.P(x, y1), R.P(x, y + h)))
                    R.cmds.append("\\draw[fill=white] %s rectangle %s;" % (R.P(lo, y0), R.P(hi, y1)))
                    R.cmds.append("\\node at %s {$\\scriptstyle %s$};" % (R.P((lo + hi) / 2, (y0 + y1) / 2), it[4]))
                else:
                    cx = sum(allx) / len(allx)
                    cy = y + 0.5 * h
                    for x in ix:
                        R.cmds.append("\\draw %s .. controls %s .. %s;" % (R.P(x, y), R.P(x, y + 0.3 * h), R.P(cx, cy)))
                    for x in ox:
                        R.cmds.append("\\draw %s .. controls %s .. %s;" % (R.P(cx, cy), R.P(x, y + 0.7 * h), R.P(x, y + h)))
                    R.cmds.append("\\draw[fill=white] %s circle (2.2pt);" % R.P(cx, cy))
            else:
                raise ValueError(kind)
        y += h
        cur = outs
    return R.cmds, y


def tikz(m, layers, scale=1.0, **kw):
    cmds, H = render(m, layers, **kw)
    body = "\n".join(cmds)
    return ("\\begin{tikzpicture}[x=%scm,y=%scm,baseline={([yshift=-0.55ex]current bounding box.center)},line width=0.6pt]\n"
            % (fmt(scale), fmt(scale)) + body + "\n\\end{tikzpicture}")


# ----------------------------------------------------------------------
# evaluation in the S_t model (boxes substituted by given S_t morphisms)
def evaluate(S, m, layers, boxes=None):
    boxes = boxes or {}
    cur = S.id(m)
    n = m
    for layer in layers:
        items = sorted(layer, key=lambda it: (it[1], arity(it)[0] > 0))
        f = S.C.mor(0, 0, {(): S.one})
        k = 0
        for it in items:
            while k < it[1]:
                f = f @ S.P
                k += 1
            nin, nout = arity(it)
            kind = it[0]
            if kind == "X":
                g = S.c
            elif kind == "T":
                g = S.tau
            elif kind == "Ts":
                g = S.taus
            elif kind == "cup":
                g = S.cup
            elif kind == "cap":
                g = S.cap
            elif kind == "box":
                g = boxes[it[4]]
            elif kind == "node":
                g = boxes.get(("node", nin, nout))
            f = f @ g
            k += nin
        while k < n:
            f = f @ S.P
            k += 1
        assert f.m == n
        cur = f * cur
        n = f.n
    return cur


# ----------------------------------------------------------------------
# conversion to planar maps (boxes are non-rotation-invariant vertices)
def to_pmap(m, layers):
    import itertools
    from pmap import PMap, smooth
    cnt = itertools.count()
    alpha, nodes, wire = {}, {}, {}
    bottom = [next(cnt) for _ in range(m)]
    cur = list(bottom)
    vid = itertools.count()
    order = {"X": [("in", 0), ("in", 1), ("out", 1), ("out", 0)],
             "T": [("in", 0), ("in", 1), ("out", 0)],
             "Ts": [("in", 0), ("out", 1), ("out", 0)]}
    for layer in layers:
        items = sorted(layer, key=lambda it: (it[1], arity(it)[0] > 0), reverse=True)
        # process right to left so indices stay valid
        for it in items:
            kind, pos = it[0], it[1]
            nin, nout = arity(it)
            ins = cur[pos:pos + nin]
            if kind == "cup":
                a, b = next(cnt), next(cnt)
                wire[a] = b; wire[b] = a
                cur = cur[:pos] + [a, b] + cur[pos:]
                continue
            if kind == "cap":
                a, b = next(cnt), next(cnt)
                wire[a] = b; wire[b] = a
                for x, y in zip(ins, (a, b)):
                    alpha[x] = y; alpha[y] = x
                cur = cur[:pos] + cur[pos + 2:]
                continue
            ind = [next(cnt) for _ in range(nin)]
            outd = [next(cnt) for _ in range(nout)]
            for x, y in zip(ins, ind):
                alpha[x] = y; alpha[y] = x
            if kind in order:
                ds = [ind[i] if s == "in" else outd[i] for s, i in order[kind]]
                tp = "T" if kind in ("T", "Ts") else "X"
            else:
                # ccw order starting at bottom-left: inputs left->right, outputs right->left
                ds = ind + list(reversed(outd))
                tp = ("BOX", it[4]) if kind == "box" else ("NODE",)
            nodes[next(vid)] = (tp, ds)
            cur = cur[:pos] + outd + cur[pos + nin:]
    n = len(cur)
    top = [next(cnt) for _ in range(n)]
    for x, y in zip(cur, top):
        alpha[x] = y; alpha[y] = x
    M = smooth(nodes, bottom + list(reversed(top)), alpha, wire)
    return M, n


def code_with_boxes(M):
    """canonical code rooted at boundary position 0 that also records, for
    darts at BOX vertices, their index (boxes are not rotation invariant)."""
    from pmap import ClosedComponent
    root = M.bdry[0]
    lab = {root: 0}
    order = [root]
    i = 0
    while i < len(order):
        d = order[i]; i += 1
        for e in (M.sigma(d), M.alpha[d]):
            if e not in lab:
                lab[e] = len(order); order.append(e)
    if len(order) != len(M.alpha):
        raise ClosedComponent()
    vof = M.vof()
    out = []
    for d in order:
        v, idx = vof[d]
        if v == "B":
            tp = ("B",)
            extra = None
        else:
            tp = M.nodes[v][0]
            extra = idx if isinstance(tp, tuple) and tp[0] == "BOX" else None
        out.append((str(tp), extra, lab[M.sigma(d)], lab[M.alpha[d]]))
    return (M.k, M.loops, tuple(out))


def isotopic(d1, d2):
    """d = (m, layers); same planar diagram with the same boundary?"""
    M1, n1 = to_pmap(*d1)
    M2, n2 = to_pmap(*d2)
    if (d1[0], n1) != (d2[0], n2):
        return False
    return code_with_boxes(M1) == code_with_boxes(M2)