#!/usr/bin/env python3
"""Aligned pictures: #2 (A1 = 0RF) vs SOMDCv6 TM7 (A1 = 0RC), same start S(r), same scale.

Writes img/eat_pair.png   one eating pass each (start of the sweep -> the eater reads a 1-digit)
       img/sweep_pair.png one full sweep each (start -> next S config, wall 1 cell further left)
       img/sweep_ours.html, img/sweep_tm7.html  tools/diagram.py viewers of the two sweeps (same start)
and prints the event numbers used in observe.json / notes.md (passes, stops, splice check).

S(r) = 0^inf 11 D> r 0^inf, digits L = 110, 0 = 000, 1 = 100. Start: r = L^4 1L1 (constructed; it is small,
both eaters eat the same 4 L's, and a drain pass re-enters an L in both machines).

Run:  python3 -B scripts/compare_tm7.py            (needs node + playwright-core + chrome-headless-shell)
      python3 -B scripts/compare_tm7.py --no-png   (numbers only)
"""
import base64
import glob
import html
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path("[local working file]")
sys.path.insert(0, str(ROOT / "tools"))
import diagram as dg  # noqa: E402
import sim  # noqa: E402

MDIR = Path(__file__).resolve().parent.parent
OURS = "1RB0RF_1LC0RE_0LC0LD_1RD1LB_0RF1RA_0RB---"
TM7 = "1RB0RC_1LC0RE_0LC0LD_1RD1LB_0RF1RA_0RB---"
R0 = "LLLL1L1"
INV = {"L": "110", "0": "000", "1": "100"}
CFG = "0^inf 11 D> " + " ".join(INV[d] for d in R0) + " 0^inf"
SUP = str.maketrans("0123456789", "⁰¹²³⁴⁵⁶⁷⁸⁹")


def word(r):
    """'LLLL1L1' -> 'L⁴ 1L1' (leading L-run as a power)."""
    n = len(r) - len(r.lstrip("L"))
    head = ("L" + str(n).translate(SUP) if n > 1 else "L" * n)
    return (head + " " + r[n:]).strip()


def cell(row, x):
    i = x - row.clo
    return row.cells[i] if 0 <= i < len(row.cells) else 0


def sweep_rows(tm):
    """Rows 0..n of one sweep from S(R0): n = first step with D at cell -1 and the S shape (wall at -3,-2)."""
    rows, _ = dg.rows_all(tm, sim.Config.parse(CFG), 5000)
    for i, r in enumerate(rows):
        if i and r.state == "D" and r.pos == -1 and cell(r, -2) == 1 and cell(r, -3) == 1 and \
                all(cell(r, k) == 0 for k in range(r.lo, -3)):
            return rows[:i + 1]
    raise RuntimeError("no sweep end")


def digits_from(row, start):
    hi = max([k for k in range(row.lo, row.hi + 1) if cell(row, k)] or [start])
    s = "".join(str(cell(row, k)) for k in range(start, hi + 1))
    s += "0" * (-len(s) % 3)
    D = {v: k for k, v in INV.items()}
    return "".join(D.get(s[i:i + 3], "?") for i in range(0, len(s), 3))


def analyse(rows):
    a1 = [i for i, r in enumerate(rows[:-1]) if r.state == "A" and cell(r, r.pos) == 1]
    stops = [i for i, r in enumerate(rows) if r.state == "D" and r.pos >= 0 and r.pos % 3 == 0
             and cell(r, r.pos) == 1 and cell(r, r.pos + 1) == 0 and cell(r, r.pos + 2) == 0]
    return a1, stops


def passes_of(rows, a1):
    """Ours: each A1 at row s (head x) starts a +1 pass that ends when the head is back at x+1 in state C."""
    out = []
    for s in a1:
        x = rows[s].pos
        e = next(i for i in range(s + 1, len(rows)) if rows[i].state == "C" and rows[i].pos == x + 1)
        out.append({"s": s, "e": e, "x": x, "max": max(r.pos for r in rows[s:e + 1]), "extra": e - s - 1})
    return out


def facts():
    ro, rt = sweep_rows(OURS), sweep_rows(TM7)
    a1o, sto = analyse(ro)
    a1t, stt = analyse(rt)
    P = passes_of(ro, a1o)
    # eater stops: the first stop, then the first stop after the restart A1
    s1o, s1t = sto[0], stt[0]
    rso = next(s for s in a1o if s > s1o)
    rst = next(s for s in a1t if s > s1t)
    s2o = next(s for s in sto if s > rso)
    s2t = next(s for s in stt if s > rst)
    # splice: cut each pass (rows s+1..e-1) out of ours; compare (state, pos) with TM7 row by row
    cut = {p["s"]: p["e"] for p in P}
    seq, i = [], 0
    while i < len(ro):
        seq.append(i)
        i = cut[i] if i in cut else i + 1
    same = 0
    for io, it in zip(seq, range(len(rt))):
        if (ro[io].state, ro[io].pos) != (rt[it].state, rt[it].pos):
            break
        same += 1
    # first divergence: which symbol each machine reads at the last common row
    lo_, lt_ = ro[seq[same - 1]], rt[same - 1]
    # drain passes after the last stop: right-to-left turns beyond the stop cell, minus the 1-digit conversion
    def turns(rows, after, beyond):
        """Drain passes after the last stop = right-to-left turns beyond the stop cell, minus the first one
        (a 4-step bounce that turns the 1-digit into L, not a pass)."""
        t = [i for i in range(after + 1, len(rows) - 1)
             if rows[i].dir > 0 and rows[i + 1].dir < 0 and rows[i].pos > beyond]
        assert t[0] - after <= 5, t
        return t[1:]
    f = {
        "start": CFG, "r": R0,
        "ours": {"n": len(ro) - 1, "a1": a1o, "passes": P, "stop1": s1o, "restart": rso, "stop2": s2o,
                 "stop1_cell": ro[s1o].pos, "stop2_cell": ro[s2o].pos, "end": digits_from(ro[-1], -1),
                 "turns": turns(ro, s2o, ro[s2o].pos)},
        "tm7": {"n": len(rt) - 1, "a1": a1t, "stop1": s1t, "restart": rst, "stop2": s2t,
                "stop1_cell": rt[s1t].pos, "stop2_cell": rt[s2t].pos, "end": digits_from(rt[-1], -1),
                "turns": turns(rt, s2t, rt[s2t].pos)},
        "splice": {"tm7_rows_identical": same, "ours_row": seq[same - 1], "tm7_row": same - 1,
                   "cell": lt_.pos, "tm7_reads": cell(lt_, lt_.pos), "ours_reads": cell(lo_, lo_.pos)},
    }
    return f, ro, rt


# ---------------------------------------------------------------- drawing (HTML + SVG, screenshotted)
BLUE, AMBER, TXT, MUTED, BG = "#3d9bff", "#ffb020", "#e9e9e9", "#9a9a9a", "#161616"


def png_uri(rows, c0, c1, n, kx, ky):
    img = dg.render_rows(rows[:n + 1], 2, style="default", cols=(c0, c1)).scaled(kx, ky)
    return "data:image/png;base64," + base64.b64encode(img.to_png()).decode()


def page(kind, f, ro, rt):
    """One composite figure: ours | TM7, same start, same scale, every step; notes laid out in the gutters."""
    eat = kind == "eat"
    o, m = f["ours"], f["tm7"]
    no = o["stop1"] if eat else o["n"]
    nt = m["stop1"] if eat else m["n"]
    c0, c1 = (-4, 24) if eat else (-4, 34)
    CW, RH = (6, 4) if eat else (5, 2)        # CSS px per cell / per step (the PNG is 2x: DPR)
    DPR = 2
    W = (c1 - c0 + 1) * CW
    AX, G, RM, TOP = 40, 176, 162, 70
    H = no * RH + RH
    X1, X2 = AX, AX + W + G
    TW = X2 + W + RM
    s = []

    def X(base, c):
        return base + (c - c0) * CW

    def Y(step):                               # vertical centre of row `step`
        return TOP + step * RH + RH / 2

    def t(x, y, txt, size=11, color=TXT, anchor="start", weight="normal"):
        s.append(f'<text x="{x:.1f}" y="{y:.1f}" font-size="{size}" fill="{color}" text-anchor="{anchor}" '
                 f'font-weight="{weight}">{html.escape(txt)}</text>')

    def line(x1, y1, x2, y2, color, w=1, dash=None):
        d = f' stroke-dasharray="{dash}"' if dash else ""
        s.append(f'<line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" stroke="{color}" '
                 f'stroke-width="{w}"{d}/>')

    # titles + start-tape ruler (digits over their 3 cells; the head starts on cell 0 in state D)
    t(X1 + W / 2, 16, "#2 (ours) · A1 = 0RF", 13, TXT, "middle", "bold")
    t(X2 + W / 2, 16, "SOMDCv6 TM7 · A1 = 0RC", 13, TXT, "middle", "bold")
    for base in (X1, X2):
        yR = TOP - 24
        t(base - 6, yR + 4, "start", 9, MUTED, "end")
        t(X(base, -1.0), yR + 4, "11", 10, MUTED, "middle")
        for i, d in enumerate(f["r"]):
            t(X(base, 3 * i + 1.5), yR + 4, d, 12, TXT, "middle", "bold")
        for i in range(len(f["r"]) + 1):
            line(X(base, 3 * i), yR - 8, X(base, 3 * i), TOP - 3, "#555")
        s.append(f'<path d="M{X(base, 0.5) - 4:.1f},{TOP - 10} l8,0 l-4,6 z" fill="#ff3b3b"/>')
        t(X(base, 0.5) + 6, TOP - 3, "D", 9, "#ff6b6b")
    # step axes on both panels
    tick, lab = (10, 20) if eat else (10, 50)
    for base, n in ((X1, no), (X2, nt)):
        for k in range(0, n + 1, tick):
            line(base - (6 if k % lab == 0 else 3), Y(k), base, Y(k), "#888")
            if k % lab == 0:
                t(base - 8, Y(k) + 3.5, f"{k:,}", 9, MUTED, "end")
        t(base - 8, TOP + n * RH + RH + 12, "step", 9, MUTED, "end")
    imgs = [(X1, png_uri(ro, c0, c1, no, CW * DPR, RH * DPR), no),
            (X2, png_uri(rt, c0, c1, nt, CW * DPR, RH * DPR), nt)]

    notes = {X1: [], X2: []}                   # gutter notes: (anchor y, [(text, size, colour, weight)])
    # ours: +1 pass boxes
    P = [p for p in o["passes"] if p["e"] <= no]
    for k, p in enumerate(P, 1):
        x0, x1 = X(X1, p["x"]) - 1.5, X(X1, p["max"] + 1) + 1.5
        y0, y1 = TOP + p["s"] * RH - 1, TOP + (p["e"] + 1) * RH + 1
        s.append(f'<rect x="{x0:.1f}" y="{y0:.1f}" width="{x1 - x0:.1f}" height="{y1 - y0:.1f}" fill="none" '
                 f'stroke="{BLUE}" stroke-width="2"/>')
        drain = p["s"] > o["stop1"]
        notes[X1].append((Y(p["s"]), [(f"+1 pass {k}", 10.5, BLUE, "bold")]
                          + ([("fired as a drain pass hit an L", 9.5, BLUE, "normal")] if drain else [])
                          + [(f"{p['extra']} extra steps", 9.5, BLUE, "normal")], BLUE, x1))
    # TM7: its A1s (it turns straight back: no pass)
    for s1 in [a for a in m["a1"] if a <= nt]:
        drain = s1 > m["stop1"]
        s.append(f'<path d="M{X(X2, m_pos(rt, s1)) - 9:.1f},{Y(s1):.1f} l-7,-4 l0,8 z" fill="{BLUE}"/>')
        notes[X2].append((Y(s1), [("A1: turns straight back", 9.5, BLUE, "normal")]
                          + ([("(a drain pass hit an L)", 9.5, BLUE, "normal")] if drain else []), BLUE, X2 + W))
    # eater stops (dashed amber across the panel)
    for base, d in ((X1, o), (X2, m)):
        rows_ = [("stop1", "eater stops at a 1-digit")] + ([] if eat else [("stop2", "eater stops again")])
        for key, txt in rows_:
            y = Y(d[key])
            line(base - 2, y, base + W + 2, y, AMBER, 1.5, "5,3")
            cellk = d[key + "_cell"]
            sub = (f"step {d[key]}, cell {cellk}, k = {cellk // 3}" if eat else
                   f"step {d[key]}, cell {cellk}" + (f", k = {len(d['turns'])}" if key == "stop2" else ""))
            notes[base].append((y, [(txt, 10.5, AMBER, "bold"), (sub, 9.5, AMBER, "normal")], AMBER, base + W))
    if not eat:
        for base, d, n in ((X1, o, no), (X2, m, nt)):
            y = Y(n)
            notes[base].append((y, [("sweep ends", 10.5, AMBER, "bold"), (f"step {n}", 9.5, AMBER, "normal")],
                                AMBER, base + W))
    # lay out each gutter: keep order, push notes down so they never overlap, leader line to the anchor row
    for base, items in notes.items():
        items.sort(key=lambda it: it[0])
        xt = base + W + 14
        nxt = -1e9
        for ay, lines, col, ax_ in items:
            h = 13 * len(lines)
            ty = max(ay - 4, nxt)
            line(ax_ + 1, ay, xt - 3, ty + 0, col, 1, "2,2")
            for j, (txt, size, colr, wt) in enumerate(lines):
                t(xt, ty + 4 + 13 * j, txt, size, colr, weight=wt)
            nxt = ty + h + 4
    if not eat:                                # drain bracket
        for base, d, n in ((X1, o, no), (X2, m, nt)):
            y0, y1 = Y(d["stop2"]) + 30, Y(n) - 30
            xb = base + W + 12
            s.append(f'<path d="M{xb - 4},{y0:.1f} l4,0 l0,{y1 - y0:.1f} l-4,0" fill="none" stroke="{TXT}"/>')
            t(xb + 6, (y0 + y1) / 2, f"drain: {len(d['turns'])} passes", 10.5, TXT, weight="bold")
            t(xb + 6, (y0 + y1) / 2 + 13, "(no A1: same", 9.5, MUTED)
            t(xb + 6, (y0 + y1) / 2 + 25, "transitions in both)", 9.5, MUTED)
        for base, d in ((X1, o), (X2, m)):    # end tape ruler: the wall is one cell further left
            yR = TOP + H + 16
            t(X(base, -2.5), yR, "11", 10, MUTED, "middle")
            for i, ch in enumerate(d["end"]):
                t(X(base, -1 + 3 * i + 1.5), yR, ch, 11, TXT, "middle", "bold")
            for i in range(len(d["end"]) + 1):
                xx = X(base, -1 + 3 * i)
                line(xx, TOP + H + 4, xx, yR + 4, "#555")
            t(base + W / 2, yR + 17, f"end: S({word(d['end'])}), wall 1 cell left", 10, MUTED, "middle")
    # footer
    sp = f["splice"]
    if eat:
        cut = sum(p["extra"] for p in P)
        L = [(f"Cut ours' {len(P)} blue passes ({cut} extra steps): {no} − {cut} = {nt}. What is left is TM7's run, "
              f"step for step", TXT, 10.5),
             ("(same state, same head cell). Each pass only adds 1 to the digits right of it.", TXT, 10.5)]
    else:
        L = [(f"Cut ours' {len(P)} blue passes: the head path equals TM7's up to TM7 step {sp['tm7_row']}. There TM7 "
              f"reads a 1 at cell {sp['cell']};", TXT, 10.5),
             (f"ours' +1s had made that digit 0, so ours eats 2 more digits (k = {len(o['turns'])} vs "
              f"{len(m['turns'])}) and drains {len(o['turns'])} passes vs {len(m['turns'])}.", TXT, 10.5)]
    L += [(f"Start (constructed, same for both): S({word(f['r'])}) = {f['start']}", MUTED, 9.5),
          (f"Rows: every step, 0..{no} (ours) and 0..{nt} (TM7), same scale.  Window: cells {c0}..{c1}.", MUTED, 9.5),
          ("white = 1, black = 0, red/green = head after a move R/L (start row: no head colour).", MUTED, 9.5),
          ("Blue box = ours' extra +1 pass (fired by A1 = 0RF).  Blue ▶ = TM7's A1 (no pass).  Dashed amber = eater stops.",
           MUTED, 9.5)]
    yF = TOP + H + (30 if eat else 58)
    for j, (txt, col, size) in enumerate(L):
        t(12, yF + 15 * j, txt, size, col)
    TH = int(yF + 15 * len(L) + 4)
    svg = "\n".join(s)
    im = "\n".join(f'<img src="/busy-beaver/machines/1RB0RF_1LC0RE_0LC0LD_1RD1LB_0RF1RA_0RB---/scripts/</span><span class="si">{u}" style="position:absolute;left:{x}px;top:{TOP}px;width:{W}px;'
                   f'height:{n * RH + RH}px;image-rendering:pixelated">' for x, u, n in imgs)
    return TW, TH, DPR, f"""<!doctype html><html><head><meta charset="utf-8"><title>{kind}</title>
<style>html,body{{margin:0;background:{BG}}}#fig{{position:relative;width:{TW}px;height:{TH}px;background:{BG};
font-family:'DejaVu Sans',Lato,sans-serif}}svg{{position:absolute;left:0;top:0}}</style></head><body>
<div id="fig">{im}<svg width="{TW}" height="{TH}" viewBox="0 0 {TW} {TH}" font-family="DejaVu Sans, Lato, sans-serif">
{svg}</svg></div></body></html>"""


def m_pos(rows, i):
    return rows[i].pos


SHOT_JS = r"""
const { chromium } = require(process.argv[2]);
(async () => {
  const b = await chromium.launch({ executablePath: process.argv[3], headless: true });
  const jobs = JSON.parse(process.argv[4]);
  for (const j of jobs) {
    const p = await b.newPage({ viewport: { width: j.w, height: j.h }, deviceScaleFactor: j.dpr });
    await p.goto('file://' + j.html, { waitUntil: 'load' });
    await (await p.$('#fig')).screenshot({ path: j.png });
    await p.close();
  }
  await b.close();
})().catch((e) => { console.error(e); process.exit(1); });
"""


def screenshot(jobs):
    pw = "[local working file]"
    exe = sorted(glob.glob(str(Path.home() / ".cache/ms-playwright/chromium_headless_shell-*/"
                               "chrome-headless-shell-linux64/chrome-headless-shell")))[-1]
    env = dict(os.environ)
    libs = Path("[local working file]")
    if (libs / "libnspr4.so").exists():
        env["LD_LIBRARY_PATH"] = str(libs) + ":" + env.get("LD_LIBRARY_PATH", "")
    with tempfile.TemporaryDirectory() as td:
        js = Path(td) / "shot.js"
        js.write_text(SHOT_JS)
        subprocess.run(["node", str(js), pw, exe, json.dumps(jobs)], check=True, env=env)


def main():
    f, ro, rt = facts()
    show = {k: v for k, v in f.items()}
    print(json.dumps(show, indent=1, default=str))
    if "--no-png" in sys.argv:
        return
    jobs = []
    tmp = Path(tempfile.mkdtemp(prefix="cmp_tm7_"))
    for kind in ("eat", "sweep"):
        w, h, dpr, doc = page(kind, f, ro, rt)
        hp = tmp / f"{kind}_pair.html"
        hp.write_text(doc)
        jobs.append({"html": str(hp), "png": str(MDIR / "img" / f"{kind}_pair.png"), "w": w, "h": h, "dpr": dpr})
    screenshot(jobs)
    for j in jobs:
        print("wrote", j["png"])
    # zoomable bbchallenge-style viewers of the two sweeps (hover shows step, state, cell)
    for tm, name, n in ((OURS, "sweep_ours", f["ours"]["n"]), (TM7, "sweep_tm7", f["tm7"]["n"])):
        out = MDIR / "img" / f"{name}.html"
        subprocess.run([sys.executable, "-B", str(ROOT / "tools/diagram.py"), tm, "--from-config", CFG,
                        "--steps", str(n), "--blaze-steps", str(n), "--records-steps", str(n),
                        "--state-steps", str(n), "-o", str(out)], check=True, stdout=subprocess.DEVNULL)
        print("wrote", out)


if __name__ == "__main__":
    main()