"""Check that every pictured identity in the note is correct."""
import sys, random
import os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fractions import Fraction as F
import figs
from figs import D, PD
from diag import evaluate, isotopic, to_pmap, code_with_boxes
from part import to_cyclic_blocks
from st import St
from deligne import from_cyclic, to_cyclic, rotate
from evalmap import sfp
from relations import build_relations
from gauge import gauge_errors, plus_eigen_basis
from system import err_offsets, ERRDIMS

t = F(37, 10)
S = St(t)
ok = True
def check(name, cond):
    global ok
    ok &= bool(cond)
    print(("  ok   " if cond else "  FAIL ") + name)

def cyc(f):
    res = {}
    for p_, v in f.d.items():
        res[to_cyclic(p_, f.m, f.n)] = v
    return [res.get(q, 0) for q in sfp(f.m + f.n)]

def mor(m, n, vec):
    f = S.C.mor(m, n, {})
    for q, x in zip(sfp(m + n), vec):
        if x:
            f = f + (S.id(n) * S.C.mor(m, n, {from_cyclic(q, m, n): S.one}) * S.id(m)).scale(x)
    return f

def basis(m, n, blocks):
    q = to_cyclic_blocks(m, n, blocks)
    return S.id(n) * S.C.mor(m, n, {from_cyclic(q, m, n): S.one}) * S.id(m)

def nodeval(nin, nout):
    q = from_cyclic(tuple([tuple(range(nin + nout))]), nin, nout)
    return S.id(nout) * S.C.mor(nin, nout, {q: S.one}) * S.id(nin)

def ev(name, boxes=None):
    m, layers = D[name]
    bx = dict(boxes or {})
    bx[("node", 2, 2)] = nodeval(2, 2)
    return evaluate(S, m, layers, bx)

def pev(name):
    m, n, blocks, _ = PD[name]
    return basis(m, n, blocks)

def eq(f, g):
    return cyc(f) == cyc(g)

print("-- partition pictures are honest: drawn diagram (tau vertices, crossings, nodes) = [pi]")
# evaluate the partition picture as a diagram: blocks of size 3 are tau, crossings c; by symmetry this is [pi]
# (checked by comparing selected word diagrams below)
check("psi_id = Pid5", eq(ev("psi_id"), pev("Pid5")))
check("psi_e = Q", eq(ev("psi_e"), pev("Q")))
check("psi_c = Pc5", eq(ev("psi_c"), pev("Pc5")))
check("psi_nu = Pnu5 - t^-1 Pcut5", eq(ev("psi_nu"), pev("Pnu5") - pev("Pcut5").scale(1 / t)))
check("I = nu4 - t^-1 e", eq(ev("I"), nodeval(2, 2) - ev("e").scale(1 / t)))
check("H = nu4 - t^-1 id", eq(ev("H"), nodeval(2, 2) - S.id(2).scale(1 / t)))
check("tau = [b1b2t1], taus = ptaus, cross = pc, e = pe", eq(ev("tau"), basis(2, 1, [["b1", "b2", "t1"]])) and eq(ev("taus"), pev("ptaus")) and eq(ev("cross"), pev("pc")) and eq(ev("e"), pev("pe")))
check("relations: loop, lollipop, bigon, I-H", cyc(ev("loop")) == [t - 1] and all(x == 0 for x in cyc(ev("lollipop"))) and eq(ev("bigon"), S.P.scale((t - 2) / t)) and eq(ev("I") - ev("H"), (S.id(2) - ev("e")).scale(1 / t)))
check("relations: curl, twist, R2, R2.5, R3", eq(ev("curl"), S.P) and eq(ev("twist"), ev("tau")) and eq(ev("Rtwo"), S.id(2)) and eq(ev("Lam"), ev("Rho")) and eq(ev("Tri"), ev("rTri")))

print("-- isotopies")
rho3 = lambda name: None
def rot_word(m, layers, k):
    """rho^k as a word: rho(f) = [cup at m, f+1, cap at 0]"""
    M, n = to_pmap(m, layers)
    for _ in range(k):
        layers = [[("cup", m)]] + [[(it[0], it[1] + 1) + tuple(it[2:]) for it in Lr] for Lr in layers] + [[("cap", 0)]]
    return (m, layers)
check("XTX = rho^3(Lam)", isotopic(rot_word(*D["Lam"], 3), D["XTX"]))
check("XT = rho^3(Rho)", isotopic(rot_word(*D["Rho"], 3), D["XT"]))
def rotcode(name, k):
    return code_with_boxes(to_pmap(*D[name])[0].rotated(k))
def code(name):
    return code_with_boxes(to_pmap(*D[name])[0])
check("Ttaus1 = rho(Lam), Ttaus2 = rho(Rho)", rotcode("Lam", 1) == code("Ttaus1") and rotcode("Rho", 1) == code("Ttaus2"))
check("rTri = rho(Tri), Tri = rho^2(Tri)", isotopic(rot_word(*D["Tri"], 1), D["rTri"]) and isotopic(rot_word(*D["Tri"], 2), D["Tri"]))
check("rhoE3 = rho(E3) (one click)", rotcode("Ebarthree", 1) == code("rhoE3"))
check("rhof = rho(f)", code_with_boxes(to_pmap(*D["fbox"])[0].rotated(1)) == code_with_boxes(to_pmap(*D["rhof"])[0]))
for a, b in (("natA_l", "natA_r"), ("natB_l", "natB_r"), ("natC_l", "natC_r"), ("natD_l", "natD_r")):
    check(f"{a} = rho^k Lam and {b} = rho^k Rho (same k)", any(rotcode("Lam", k) == code(a) and rotcode("Rho", k) == code(b) for k in range(5)))
for a, b in (("natcap_l", "natcap_r"), ("natcup_l", "natcup_r")):
    check(f"{a} = rho^k R2 and {b} = rho^k id", any(rotcode("Rtwo", k) == code(a) and rotcode("strandtwo", k) == code(b) for k in range(4)))
check("curlcap, crosscup are rotations of curl", isotopic(D["curlcap"], (2, [[("X", 0)], [("cap", 0)]])))
# Etfrot: hand drawn; its word form is [cup0, cup1, box at 2, cap3, cap2]
Erot = (3, [[("cup", 0)], [("cup", 1)], [("box", 2, 3, 2, "E")], [("cap", 3)], [("cap", 2)]])
check("hand-drawn Etfrot is rho^3(E25) (via its word form)", code_with_boxes(to_pmap(3, [[("box", 0, 3, 2, "E")]])[0].rotated(3)) == code_with_boxes(to_pmap(*Erot)[0]))

print("-- pictured overlap identities hold for the genuine first-order deformations c -> c + eps w")
rels = build_relations(t)
off, _ = err_offsets()
random.seed(1)
for w in plus_eigen_basis(t):
    g = gauge_errors(rels, w, t)
    a = g[off["a"]]; b = g[off["b"]]
    E2 = mor(2, 2, g[off["E2"]:off["E2"] + 4])
    E25 = mor(3, 2, g[off["E25"]:off["E25"] + 11])
    E3 = mor(3, 3, g[off["E3"]:off["E3"] + 41])
    bx = {r"\overline{E}_2": E2, r"\overline{E}_{2.5}": E25, r"\overline{E}_3": E3}
    check(" cap E2 = 2a cap", eq(ev("capE", bx), S.cap.scale(2 * a)))
    check(" tau E2 = 2b tau", eq(ev("tauE", bx), S.tau.scale(2 * b)))
    E25rot3 = mor(3, 2, [None] * 0 or cyc(rotate(E25, 3, 2, 3)))
    check(" level5: CE25 + rho^3 E25 = E2XT", eq(ev("CE25", bx) + E25rot3, ev("E2XT", bx)))
    check(" cup E25 = a tau*", eq(ev("cupE25", bx), S.taus.scale(a)))
    check(" level6: c1E3 + E3c2 = XXE2 - E2XX", eq(ev("E3c", bx) + ev("cE3", bx), ev("XXE2", bx) - ev("E2XX", bx)))
    E25rot1 = rotate(E25, 2, 3, 1)
    check(" level6 merge: E3(tau* x 1) = rho(E25) - c2 rho(E25)", eq(ev("E3taus", bx), E25rot1 - (S.P @ S.c) * E25rot1))
    check(" rho E3 = -E3", eq(rotate(E3, 3, 3, 1), E3.scale(-1)))

print("-- combinatorial facts")
for i, (s, bl) in enumerate(figs.M6):
    pass
m6 = sum((basis(3, 3, bl).scale(s) for s, bl in figs.M6[1:]), basis(3, 3, figs.M6[0][1]).scale(figs.M6[0][0]))
mer = sum((basis(2, 3, bl).scale(s) for s, bl, _ in figs.MER[1:]), basis(2, 3, figs.MER[0][1]).scale(figs.MER[0][0]))
check("m6 o (tau* x 1) = signed sum of mer pictures", eq(m6 * (S.taus @ S.P), mer))
check("rho m6 = -m6 and c1 m6 c2 = -m6", eq(rotate(m6, 3, 3, 1), m6.scale(-1)) and eq((S.c @ S.P) * m6 * (S.P @ S.c), m6.scale(-1)))
k5 = [(1, "Q"), (-1, "P2"), (1, "K3"), (-1, "K4"), (1, "K5"), (-1, "P6")]
K5 = sum((pev(nm).scale(s) for s, nm in k5[1:]), pev("Q"))
check("k5 o (cup x 1) = t tau*", eq(K5 * (S.cup @ S.P), S.taus.scale(t)))
check("(sigma + rho^3) k5 = 0", eq(K5 * (S.c @ S.P) + rotate(K5, 3, 2, 3), S.C.mor(3, 2, {})))
print("ALL OK" if ok else "FAILURES")