"""All diagrams of the note, generated from layered words (so every picture is
an honest planar string diagram whose only crossings are the crossing
generator), plus checks that the pictured identities are correct."""
import sys
sys.path.insert(0, __import__("os").path.dirname(__import__("os").path.abspath(__file__)))
from fractions import Fraction as F
from diag import tikz, evaluate, isotopic, to_pmap, code_with_boxes
from part import part_tikz, to_cyclic_blocks

X = lambda i: ("X", i)
T = lambda i: ("T", i)
Ts = lambda i: ("Ts", i)
CUP = lambda i: ("cup", i)
CAP = lambda i: ("cap", i)
BOX = lambda i, a, b, l: ("box", i, a, b, l)
NODE = lambda i, a, b: ("node", i, a, b)


def L(*layers):
    return [list(l) if isinstance(l, (list, tuple)) and l and isinstance(l[0], tuple) else [l] for l in layers]


D = {}      # name -> (m, layers)
PD = {}     # name -> (m, n, blocks, vpos)
HAND = {}   # name -> raw tikz


def d(name, m, *layers):
    D[name] = (m, L(*layers))


def p(name, m, n, blocks, vpos=None):
    PD[name] = (m, n, blocks, vpos)


# ---------------- generators and basic diagrams ----------------
d("strand", 1)
d("strandtwo", 2)
d("strandthree", 3)
d("cup", 0, CUP(0))
d("cap", 2, CAP(0))
d("tau", 2, T(0))
d("taus", 1, Ts(0))
d("cross", 2, X(0))
d("e", 2, CAP(0), CUP(0))
d("I", 2, T(0), Ts(0))
d("H", 2, Ts(0), T(1))
d("loop", 0, CUP(0), CAP(0))
d("lollipop", 1, Ts(0), CAP(0))
d("bigon", 1, Ts(0), T(0))
d("node4", 2, NODE(0, 2, 2))
# rotation invariance pictures: rho(f) = [cup at m, f at 1, cap at 0]
d("taurot", 2, CUP(2), T(1), CAP(0))
d("crossrot", 2, CUP(2), X(1), CAP(0))
d("rhof", 2, CUP(2), BOX(1, 2, 2, "f"), CAP(0))
d("fbox", 2, BOX(0, 2, 2, "f"))
# relations
d("curl", 1, CUP(1), X(0), CAP(1))
d("twist", 2, X(0), T(0))
d("Rtwo", 2, X(0), X(0))
# zigzags (monoidal deformations are pivotal)
d("snakeL", 1, CUP(1), CAP(0))
d("snakeR", 1, CUP(0), CAP(1))
d("zzL", 2, CUP(1), CAP(0), CAP(0))
d("zzR", 2, CUP(1), CAP(2), CAP(0))
d("Lam", 3, X(0), X(1), T(0))
d("Rho", 3, T(1), X(0))
d("Tri", 3, X(0), X(1), X(0))
d("rTri", 3, X(1), X(0), X(1))
d("Etwo", 2, BOX(0, 2, 2, r"E_2"))
d("Etf", 3, BOX(0, 3, 2, r"E_{2.5}"))
d("Ethree", 3, BOX(0, 3, 3, r"E_3"))
d("wbox", 2, BOX(0, 2, 2, "w"))
d("Ebar2", 2, BOX(0, 2, 2, r"\overline{E}_2"))
d("Ebartf", 3, BOX(0, 3, 2, r"\overline{E}_{2.5}"))
d("Ebarthree", 3, BOX(0, 3, 3, r"\overline{E}_3"))
# rotated forms
d("curlcap", 2, X(0), CAP(0))
d("crosscup", 0, CUP(0), X(0))
d("twistdown", 1, Ts(0), X(0))
# naturality forms (Lemma: presentation)
d("natA_l", 3, X(0), X(1), T(0))
d("natA_r", 3, T(1), X(0))
d("natB_l", 3, X(1), X(0), T(1))
d("natB_r", 3, T(0), X(0))
d("natC_l", 2, Ts(1), X(0), X(1))
d("natC_r", 2, X(0), Ts(0))
d("natD_l", 2, Ts(0), X(1), X(0))
d("natD_r", 2, X(0), Ts(1))
d("natcap_l", 3, X(0), X(1), CAP(0))
d("natcap_r", 3, CAP(1))
d("natcup_l", 1, CUP(0), X(1), X(0))
d("natcup_r", 1, CUP(1))

# ---------------- level 4 ----------------
d("capcc", 2, X(0), X(0), CAP(0))
d("capE", 2, BOX(0, 2, 2, r"\overline{E}_2"), CAP(0))
d("taucc", 2, X(0), X(0), T(0))
d("tauE", 2, BOX(0, 2, 2, r"\overline{E}_2"), T(0))
d("psix", 3, BOX(0, 2, 2, "x"), X(1), T(0))
d("psiE", 3, BOX(0, 2, 2, r"\overline{E}_2"), X(1), T(0))
d("psi_id", 3, X(1), T(0))
d("psi_e", 3, CAP(0), CUP(0), X(1), T(0))
d("psi_c", 3, X(0), X(1), T(0))
d("psi_nu", 3, NODE(0, 2, 2), X(1), T(0))
d("wc", 2, X(0), BOX(0, 2, 2, "w"))
d("cw", 2, BOX(0, 2, 2, "w"), X(0))

# ---------------- level 5 ----------------
d("LamC", 3, X(0), X(0), X(1), T(0))
d("XT", 3, X(1), T(0))
d("E2XT", 3, BOX(0, 2, 2, r"\overline{E}_2"), X(1), T(0))
d("XTX", 3, X(0), T(1), X(0))
d("CE25", 3, X(0), BOX(0, 3, 2, r"\overline{E}_{2.5}"))
d("cupLam", 1, CUP(0), X(0), X(1), T(0))
d("cupXT", 1, CUP(0), X(1), T(0))
d("cupRho", 1, CUP(0), T(1), X(0))
d("cupE25", 1, CUP(0), BOX(0, 3, 2, r"\overline{E}_{2.5}"))
d("cupk", 1, CUP(0), BOX(0, 3, 2, r"k_5"))
d("kbox", 3, BOX(0, 3, 2, r"k_5"))

# ---------------- level 6 ----------------
d("cT", 3, X(0), X(1), X(0), X(0))
d("XX", 3, X(0), X(1))
d("XXE2", 3, X(0), X(1), BOX(0, 2, 2, r"\overline{E}_2"))
d("cTalt", 3, X(1), X(0), X(1), X(0))
d("E3c", 3, BOX(0, 3, 3, r"\overline{E}_3"), X(0))
d("cE3", 3, X(1), BOX(0, 3, 3, r"\overline{E}_3"))
d("XXXX2", 3, X(1), X(1), X(0), X(1))
d("E2XX", 3, BOX(1, 2, 2, r"\overline{E}_2"), X(0), X(1))
d("Ttaus", 2, Ts(0), X(0), X(1), X(0))
d("Ttaus1", 2, Ts(0), X(1), X(0))
d("Ttaus2", 2, X(0), Ts(1))
d("rTtaus", 2, Ts(0), X(1), X(0), X(1))
d("rTtaus1", 2, X(0), Ts(1), X(1))
d("E3taus", 2, Ts(0), BOX(0, 3, 3, r"\overline{E}_3"))
d("mtaus", 2, Ts(0), BOX(0, 3, 3, r"m_6"))
d("mbox", 3, BOX(0, 3, 3, r"m_6"))
d("rhoE3", 3, CUP(3), BOX(1, 3, 3, r"\overline{E}_3"), CAP(0))

# ---------------- partition diagrams ----------------
# 4-boxes (m=n=2)
p("pid", 2, 2, [["b1", "t1"], ["b2", "t2"]])
p("pe", 2, 2, [["b1", "b2"], ["t1", "t2"]])
p("pc", 2, 2, [["b1", "t2"], ["b2", "t1"]])
p("pnu", 2, 2, [["b1", "b2", "t1", "t2"]])
# 5-boxes (m=3, n=2)
p("Q", 3, 2, [["b1", "b2"], ["b3", "t1", "t2"]])
p("P2", 3, 2, [["b1", "b2", "t1"], ["b3", "t2"]])
p("P3", 3, 2, [["b1", "b3", "t2"], ["b2", "t1"]], vpos={0: (1.55, 0.5)})
p("P4", 3, 2, [["b1", "t2"], ["b2", "b3", "t1"]], vpos={1: (1.45, 0.38)})
p("P5", 3, 2, [["b1", "t1", "t2"], ["b2", "b3"]])
p("P6", 3, 2, [["b1", "b2", "b3"], ["t1", "t2"]], vpos={0: (1.0, 0.42)})
p("K3", 3, 2, [["b1", "t1"], ["b2", "b3", "t2"]])
p("K4", 3, 2, [["b1", "b3", "t1"], ["b2", "t2"]], vpos={0: (0.55, 0.62)})
p("K5", 3, 2, [["b1", "b3"], ["b2", "t1", "t2"]])
p("Pid5", 3, 2, [["b1", "b3", "t1"], ["b2", "t2"]], vpos={0: (0.55, 0.62)})
p("Pc5", 3, 2, [["b1", "t2"], ["b2", "b3", "t1"]], vpos={1: (1.45, 0.38)})
p("Pnu5", 3, 2, [["b1", "b2", "b3", "t1", "t2"]])
p("Pcut5", 3, 2, [["b1", "b2", "t2"], ["b3", "t1"]], vpos={0: (0.6, 0.38)})
# 3-box
p("ptaus", 1, 2, [["b1", "t1", "t2"]])
# 6-boxes (m=n=3): the 8 matchings of m6
M6 = [
    (+1, [["b1", "b2"], ["b3", "t3"], ["t2", "t1"]]),
    (-1, [["b1", "b2"], ["b3", "t2"], ["t3", "t1"]]),
    (-1, [["b1", "b3"], ["b2", "t3"], ["t2", "t1"]]),
    (+1, [["b1", "b3"], ["b2", "t1"], ["t3", "t2"]]),
    (+1, [["b1", "t2"], ["b2", "b3"], ["t3", "t1"]]),
    (-1, [["b1", "t2"], ["b2", "t1"], ["b3", "t3"]]),
    (-1, [["b1", "t1"], ["b2", "b3"], ["t3", "t2"]]),
    (+1, [["b1", "t1"], ["b2", "t3"], ["b3", "t2"]]),
]
for i, (s, bl) in enumerate(M6):
    p("m%d" % i, 3, 3, bl)
# merges of m6 (m=2: bottom b1 = merged point, b2 = old b3 ; n=3)
MER = [
    (-1, [["b1", "b2", "t3"], ["t2", "t1"]], {0: (1.25, 0.42)}),
    (+1, [["b1", "b2", "t2"], ["t3", "t1"]], {0: (0.6, 0.38)}),
    (-1, [["b1", "t2", "t1"], ["b2", "t3"]], {0: (0.15, 0.55)}),
    (+1, [["b1", "t3", "t1"], ["b2", "t2"]], {0: (0.1, 0.66)}),
]
for i, (s, bl, vp) in enumerate(MER):
    p("mer%d" % i, 2, 3, bl, vp)
# V_- orbit representatives (m=n=3), drawn as 6-boxes
p("V1", 3, 3, [["b1", "b2"], ["b3", "t3"], ["t2", "t1"]])
p("V2", 3, 3, [["b1", "b2"], ["b3", "t3", "t2", "t1"]])
p("V3", 3, 3, [["b1", "b2"], ["b3", "t2"], ["t3", "t1"]])
p("V4", 3, 3, [["b1", "b2", "b3", "t2"], ["t3", "t1"]], vpos={0: (1.0, 0.42)})
p("V5", 3, 3, [["b1", "b2", "t3"], ["b3", "t2", "t1"]], vpos={0: (1.2, 0.35), 1: (1.4, 0.68)})


# ---------------- hand-drawn: the coupon rotated by three clicks -------------
def rot3_E25(label=r"\overline{E}_{2.5}", scale=0.5):
    # Hom(X^3, X^2) with the coupon turned upside down (rotation by three clicks).
    # Coupon legs on its (now upper) side, left to right: old b3, old b2, old b1;
    # on its (now lower) side, left to right: old t2, old t1.
    s = r"""\begin{tikzpicture}[x=%(sc)scm,y=%(sc)scm,baseline={([yshift=-0.55ex]current bounding box.center)},line width=0.6pt]
\draw (0,0) .. controls (0,2.7) and (0.2,3.15) .. (0.6,3.15) .. controls (0.85,3.15) and (0.95,2.9) .. (0.95,2.5);
\draw (1.5,2.5) .. controls (1.5,3.1) and (0.8,3.2) .. (0.8,3.9);
\draw (2.05,2.5) .. controls (2.05,3.1) and (1.8,3.2) .. (1.8,3.9);
\draw (1,0) .. controls (1,0.8) and (1.2,0.9) .. (1.2,1.4);
\draw (2,0) .. controls (2,0.8) and (1.8,0.9) .. (1.8,1.4);
\draw[fill=white] (0.6,1.4) rectangle (2.4,2.5);
\node[rotate=180] at (1.5,1.95) {$\scriptstyle %(lab)s$};
\end{tikzpicture}""" % dict(sc=scale, lab=label)
    return s


HAND["Etfrot"] = rot3_E25()
HAND["krot"] = rot3_E25(r"k_5")

# ---------------- generic maps on box spaces ----------------
d("Cx5", 3, X(0), BOX(0, 3, 2, "x"))
d("xcup", 1, CUP(0), BOX(0, 3, 2, "x"))
d("xbox5", 3, BOX(0, 3, 2, "x"))
d("xbox6", 3, BOX(0, 3, 3, "x"))
d("cxc", 3, X(1), BOX(0, 3, 3, "x"), X(0))
d("xtaus", 2, Ts(0), BOX(0, 3, 3, "x"))
d("xrho1", 3, CUP(3), BOX(1, 3, 3, "x"), CAP(0))
d("Etaurot", 2, Ts(0), X(1), X(0))
# Gamma-orbit representatives on 6-boxes
GORB = [
    (8, [["b1", "b2"], ["b3", "t3"], ["t2", "t1"]], None, None),
    (12, [["b1", "b2"], ["b3", "t3", "t2", "t1"]], None, r"(b_3\,t_1)"),
    (6, [["b1", "b2"], ["b3", "t1"], ["t3", "t2"]], None, r"(b_3\,t_1)"),
    (4, [["b1", "b2", "b3"], ["t3", "t2", "t1"]], None, r"(b_2\,b_3)(t_2\,t_1)"),
    (1, [["b1", "b2", "b3", "t3", "t2", "t1"]], None, r"(b_2\,t_2)"),
    (6, [["b1", "b2", "t3"], ["b3", "t2", "t1"]], {0: (1.2, 0.35), 1: (1.4, 0.68)}, r"(b_1\,t_3)"),
    (3, [["b1", "b2", "t3", "t2"], ["b3", "t1"]], {0: (1.6, 0.72)}, r"(b_2\,t_2)"),
    (1, [["b1", "t3"], ["b2", "t2"], ["b3", "t1"]], None, r"(b_2\,t_2)"),
]
for i, (sz, bl, vp, wit) in enumerate(GORB):
    p("G%d" % i, 3, 3, bl, vp)
HAND["Pproj"] = r"""\begin{tikzpicture}[x=0.5cm,y=0.5cm,baseline={([yshift=-0.55ex]current bounding box.center)},line width=0.6pt]
\draw (0,0) -- (0,0.7); \fill (0,0.7) circle (1.5pt); \draw (0,1.2) -- (0,1.9); \fill (0,1.2) circle (1.5pt);
\end{tikzpicture}"""
HAND["xrot"] = rot3_E25("x")


def write_tex(path):
    out = ["% generated by figs.py -- do not edit by hand"]
    for name, (m, layers) in D.items():
        out.append("\\expandafter\\def\\csname D@%s\\endcsname{%s}" % (name, tikz(m, layers)))
    for name, (m, n, blocks, vpos) in PD.items():
        out.append("\\expandafter\\def\\csname D@%s\\endcsname{%s}" % (name, part_tikz(m, n, blocks, vpos=vpos)))
    for name, s in HAND.items():
        out.append("\\expandafter\\def\\csname D@%s\\endcsname{%s}" % (name, s))
    open(path, "w").write("\n".join(out) + "\n")


if __name__ == "__main__":
    write_tex(sys.argv[1] if len(sys.argv) > 1 else "diagrams.tex")
    print(len(D), "word diagrams,", len(PD), "partition diagrams,", len(HAND), "hand drawn")