#!/usr/bin/env python3
"""p8v2 side-by-side renderer, for any generation.

Implements section 15 of tests/pprog/p8v2-viz-spec.md: the GENERATION PANEL, which
puts every variant of a generation over the SAME input positions so that where they
diverge is visible where it happens.

Input is the per-position dumps the compressor writes next to each model file and
that tests/pprog/acceptance-p8v2 collects into tests/pprog/gen<N>-pos/ as
<variant>.<sample>.pos.  Output is one self-contained HTML file, inline SVG only:
no external stylesheet, no script from a CDN, no image file, nothing fetched.

Standard library only, deliberately: the repository has no numpy or matplotlib and
this must run wherever the acceptance test runs.

  python3 tests/pprog/gen1-viz.py [posdir] [outfile]

The generation number is taken from the posdir name, so both generations render
with no edit to this file and no second copy of it:

  python3 tests/pprog/gen1-viz.py                       # gen1-pos  -> gen1.html
  python3 tests/pprog/gen1-viz.py tests/pprog/gen2-pos  # gen2-pos  -> gen2.html
"""

import glob
import html
import re
import os
import sys

# --------------------------------------------------------------- S 6 palette
# The categorical palette of S 6.2, used here by MEANING, not by entity: the hue of
# a cell says which level supplied the prediction at that position.
DARK = dict(surface="#1a1a19", grid="#2c2c2a", muted="#898781", ink="#ffffff",
            c1="#3987e5", c2="#d95926", c3="#199e70", warn="#fab219")
LIGHT = dict(surface="#fcfcfb", grid="#e1e0d9", muted="#898781", ink="#0b0b0b",
             c1="#2a78d6", c2="#eb6834", c3="#1baf7a", warn="#fab219")

CELL_W, CELL_H, ROW_GAP = 11, 18, 3
LEFT, TOP = 132, 46
GUTTER = 250
STRIP_H, STRIP_GAP = 12, 4        # the divergence strip and its three component strips

# ------------------------------------------------------------- the .pos format
# One line per sample position, tab separated, after two comment lines:
#   pos byte rec causal tok_fired tok_a tok_b tok_c settled s1 w1 s2 w2 conv
COLS = ["pos", "byte", "rec", "causal", "tok_fired", "tok_a", "tok_b", "tok_c",
        "settled", "s1", "w1", "s2", "w2", "conv"]


def read_pos(path):
    """Return (axes, M, rows) where rows is a list of dicts keyed by COLS."""
    axes, M, rows = "??????", 0, []
    with open(path) as f:
        for line in f:
            line = line.rstrip("\n")
            if line.startswith("# axes"):
                parts = line.split()
                axes = parts[2]
                M = int(parts[4])
                continue
            if line.startswith("#"):
                continue
            if not line.strip():
                continue
            vals = line.split("\t")
            if len(vals) < len(COLS):
                continue
            rows.append({k: int(v) for k, v in zip(COLS, vals[:len(COLS)])})
    return axes, M, rows


def ch(v):
    """S 6.6 frame-caption rule: printable bytes as the character, else hex."""
    return chr(v) if 0x21 <= v <= 0x7E else ("SP" if v == 0x20 else "%02X" % v)


def ramp(pal, hue, w):
    """Lightness by magnitude (S 6.4): the hue at 30..100% of full, w in 0..255."""
    u = max(0.0, min(1.0, w / 255.0))
    r = int(hue[1:3], 16)
    g = int(hue[3:5], 16)
    b = int(hue[5:7], 16)
    f = 0.30 + 0.70 * u
    return "#%02x%02x%02x" % (int(r * f), int(g * f), int(b * f))


def level_hue(pal, row):
    """Which level supplied the prediction here.

    class 3 (tokens)   a k=2 token fired
    class 2 (k-models) k=1 only
    class 1 (bytes)    neither: the position is recorded, so the value came from
                       the top-level memory event rather than from a prediction
    """
    if row["rec"]:
        return pal["c1"]
    return pal["c3"] if row["tok_fired"] else pal["c2"]


def panel(sample, entries, pal, width_positions):
    """One SVG panel: every variant of the generation over the same positions."""
    n = width_positions
    h = (TOP + len(entries) * (CELL_H + ROW_GAP) + 8
         + 4 * (STRIP_H + STRIP_GAP) + 16)
    w = LEFT + n * CELL_W + GUTTER
    parts = []
    out = ['<svg viewBox="0 0 %d %d" width="%d" height="%d" role="img" '
           'aria-label="the generation over sample %s">' % (w, h, w, h, html.escape(sample))]
    out.append('<rect width="%d" height="%d" fill="%s"/>' % (w, h, pal["surface"]))

    # The bytes themselves, once, across the top: this is the shared x axis.
    base = entries[0][2]
    out.append('<text x="8" y="20" fill="%s" font-size="12" font-family="monospace">'
               'sample %s &#8212; %d positions</text>' % (pal["ink"], html.escape(sample), n))
    for p in range(n):
        x = LEFT + p * CELL_W
        out.append('<text x="%d" y="38" fill="%s" font-size="8" font-family="monospace" '
                   'text-anchor="middle">%s</text>'
                   % (x + CELL_W // 2, pal["muted"], html.escape(ch(base[p]["byte"]))))

    for r, (variant, axes, rows, meta) in enumerate(entries):
        y = TOP + r * (CELL_H + ROW_GAP)
        out.append('<text x="8" y="%d" fill="%s" font-size="10" font-family="monospace">'
                   '%s</text>' % (y + 13, pal["ink"], html.escape(variant.replace("wordsv2-", ""))))
        out.append('<text x="72" y="%d" fill="%s" font-size="10" font-family="monospace">'
                   '%s</text>' % (y + 13, pal["muted"], html.escape(axes)))
        for p in range(n):
            row = rows[p]
            x = LEFT + p * CELL_W
            fill = ramp(pal, level_hue(pal, row), 255 if row["rec"] else row["w1"])
            out.append('<rect x="%d" y="%d" width="%d" height="%d" fill="%s"/>'
                       % (x, y, CELL_W - 1, CELL_H, fill))
            if row["rec"]:
                # Trace membership is a RING, never a hue (S 6.5).
                out.append('<rect x="%.1f" y="%.1f" width="%.1f" height="%.1f" fill="none" '
                           'stroke="%s" stroke-width="2"/>'
                           % (x + 1, y + 1, CELL_W - 3, CELL_H - 2, pal["c2"]))
            if row["settled"] != row["byte"]:
                # Settled argmax disagrees with the byte: shape, not hue.
                out.append('<line x1="%d" y1="%d" x2="%d" y2="%d" stroke="%s" '
                           'stroke-width="1"/>' % (x, y + CELL_H, x + CELL_W - 1, y, pal["ink"]))
        gx = LEFT + n * CELL_W + 10
        out.append('<text x="%d" y="%d" fill="%s" font-size="10" font-family="monospace">'
                   '%s</text>' % (gx, y + 13, pal["muted"], html.escape(meta)))

    # The divergence strip: a tick under every column where the variants do not all agree on
    # (recorded, level, settled argmax). This is the whole point of the panel -- but a single
    # combined strip saturates, and a saturated strip is not a list of positions worth looking at.
    # At e64 all eleven variants agree on recorded and on level at EVERY position, so the combined
    # strip is 57 of 64 ticks and every one of them is the settled argmax. Undecomposed that reads
    # as "they all differ everywhere"; decomposed it reads as the true finding, which is that at
    # this size the generation differs on the settling axis and on nothing else. So the strip is
    # drawn as the combined row (whose count is the panel's headline number, unchanged) followed by
    # one row per component. All four are the warning step: divergence is one KIND, and giving the
    # components three hues would make hue carry a second meaning, which S 6.2 forbids. They are
    # told apart by their direct labels and their own counts in the gutter, per S 6.6.
    components = [("divergence", lambda row: (row["rec"], row["tok_fired"], row["settled"])),
                  ("&#183; recorded", lambda row: row["rec"]),
                  ("&#183; level", lambda row: row["tok_fired"]),
                  ("&#183; settled", lambda row: row["settled"])]
    y0 = TOP + len(entries) * (CELL_H + ROW_GAP) + 8
    ndiv = 0
    for c, (label, key) in enumerate(components):
        y = y0 + c * (STRIP_H + STRIP_GAP)
        bar = STRIP_H if c == 0 else STRIP_H - 4
        out.append('<text x="8" y="%d" fill="%s" font-size="10" font-family="monospace">'
                   '%s</text>' % (y + 10, pal["ink"] if c == 0 else pal["muted"], label))
        count = 0
        for p in range(n):
            if len({key(rows[p]) for _, _, rows, _ in entries}) > 1:
                count += 1
                x = LEFT + p * CELL_W
                out.append('<rect x="%d" y="%d" width="%d" height="%d" fill="%s"/>'
                           % (x, y, CELL_W - 1, bar, pal["warn"]))
        if c == 0:
            ndiv = count
        out.append('<text x="%d" y="%d" fill="%s" font-size="10" font-family="monospace">'
                   '%d of %d positions</text>'
                   % (LEFT + n * CELL_W + 10, y + 10, pal["muted"], count, n))
        parts.append((label.replace("&#183; ", ""), count))
    out.append("</svg>")
    return "\n".join(out), ndiv, parts[1:]


LEGEND = """
<div class="legend">
  <span><i style="background:%(c3)s"></i>k=2 token fired</span>
  <span><i style="background:%(c2)s"></i>k=1 only</span>
  <span><i style="background:%(c1)s"></i>recorded &#8212; from the memory event</span>
  <span><i class="ring" style="border-color:%(c2)s"></i>in the trace</span>
  <span><i class="strike"></i>settled argmax &#8800; the byte</span>
  <span><i style="background:%(warn)s"></i>the variants diverge here</span>
  <span><i class="thin" style="background:%(warn)s"></i>&#8230; and on which of the three</span>
</div>
"""


def build(posdir, outfile, gen="1"):
    files = sorted(glob.glob(os.path.join(posdir, "*.pos")))
    if not files:
        sys.stderr.write("gen1-viz: no .pos dumps in %s; run tests/pprog/acceptance-p8v2 first\n"
                         % posdir)
        return 1
    bysample = {}
    for path in files:
        stem = os.path.basename(path)[:-4]
        variant, _, sample = stem.rpartition(".")
        axes, M, rows = read_pos(path)
        if not rows:
            continue
        # The gutter carries what the row's marks cannot be counted by eye, and in particular the
        # two quantities that separate the variants where the hues do not. STRIKE is how often the
        # settled argmax missed the byte. CONV is the sweep the window's argmax stopped changing at;
        # when the sample is one window (M <= W, which is the e64 case) there is a single window and
        # so a single conv for the whole run, and it is reported as the scalar it is rather than as
        # a mean over positions that all carry the same number.
        convs = {r["conv"] for r in rows}
        conv = ("conv %d" % convs.pop() if len(convs) == 1
                else "conv ~%.0f" % (sum(r["conv"] for r in rows) / len(rows)))
        meta = "M %d  rec %d  tok %d  strike %d  %s" % (
            M, sum(r["rec"] for r in rows), sum(r["tok_fired"] for r in rows),
            sum(r["settled"] != r["byte"] for r in rows), conv)
        bysample.setdefault(sample, []).append((variant, axes, rows, meta))

    # Panels in LADDER order, smallest first, not the alphabetical order of the sample names, which
    # puts e10k before e1k before e64 -- neither size order nor any other order a reader wants. The
    # smallest sample is also the only one drawn in full, so it belongs first.
    def bysize(sample):
        entries = bysample[sample]
        return (max(int(m.split()[1]) for _, _, _, m in entries), sample)

    body = []
    for sample in sorted(bysample, key=bysize):
        entries = sorted(bysample[sample])
        n = min(128, min(len(e[2]) for e in entries))
        if n <= 0:
            continue
        dark, ndiv, parts = panel(sample, entries, DARK, n)
        light, _, _ = panel(sample, entries, LIGHT, n)
        # The heading carries the decomposition too, because the combined number on its own is the
        # one that misleads: 57 of 64 at e64 reads as "they differ everywhere" when what it means is
        # "they differ on the settling axis and agree on everything else".
        body.append('<section><h2>%s <small>&#8212; %d variants, %d positions, '
                    '%d divergent (%s)</small></h2><div class="dk">%s</div>'
                    '<div class="lt">%s</div></section>'
                    % (html.escape(sample), len(entries), n, ndiv,
                       html.escape(", ".join("%s %d" % (k, v) for k, v in parts)), dark, light))

    doc = """<!doctype html>
<meta charset="utf-8">
<title>p8v2 generation %(gen)s &#8212; the variants side by side</title>
<style>
 :root { color-scheme: dark light; }
 body { background:#1a1a19; color:#fff; font:14px/1.5 system-ui,sans-serif; margin:0 auto;
        max-width:1500px; padding:24px; }
 .lt { display:none; }
 @media (prefers-color-scheme: light) {
   body { background:#fcfcfb; color:#0b0b0b; }
   .dk { display:none; } .lt { display:block; }
 }
 h1 { font-size:20px; margin:0 0 4px; } h2 { font-size:14px; margin:28px 0 6px; font-weight:600; }
 small { color:#898781; font-weight:400; }
 p { color:#898781; max-width:78ch; }
 section > div { overflow-x:auto; }
 .legend { display:flex; flex-wrap:wrap; gap:14px; margin:10px 0 18px; font-size:12px;
           color:#898781; }
 .legend i { display:inline-block; width:12px; height:12px; margin-right:5px;
             vertical-align:-2px; }
 .legend i.ring { background:none; border:2px solid; }
 .legend i.thin { height:7px; vertical-align:0; }
 .legend i.strike { background:linear-gradient(to top right, transparent 45%%,
                    currentColor 45%%, currentColor 55%%, transparent 55%%); }
 .prov { border-left:3px solid #fab219; padding:2px 0 2px 12px; margin:18px 0; }
 code { font-family:ui-monospace,monospace; }
</style>
<h1>p8v2 generation %(gen)s &#8212; the variants side by side</h1>
<p>Section 15 of <code>tests/pprog/p8v2-viz-spec.md</code>. Every variant of the generation over the
same input positions, so that where they diverge is visible <em>where it happens</em>. Rendered from
the per-position dumps the compressor writes beside each model file; nothing here is a placeholder
model. The summary TSV in <code>tests/pprog/gen%(gen)s.tsv</code> is the weaker view of the same runs.</p>
%(legend)s
<div class="prov"><strong>Still provisional.</strong> The settling rule <em>f</em>, the keep/prune
score in <em>&omega;</em>, and the sparsification predicate are what this generation is exploring:
the marks below are what each candidate answer actually did, not what the design says it should do.
What is no longer provisional, and was in the standalone rendering: strengths are real LSA, the
composite prediction is the implemented one, and trace membership is whatever each variant's
&omega; alternative decided.</div>
%(body)s
""" % dict(legend=LEGEND % DARK, body="\n".join(body), gen=html.escape(gen))
    with open(outfile, "w") as f:
        f.write(doc)
    sys.stderr.write("gen1-viz: wrote %s (%d samples, %d dumps)\n"
                     % (outfile, len(bysample), len(files)))
    return 0


if __name__ == "__main__":
    here = os.path.dirname(os.path.abspath(__file__))
    # The generation number is read off the posdir name, so rendering generation 2 is
    # `gen1-viz.py tests/pprog/gen2-pos tests/pprog/gen2.html` and nothing in here is edited.
    posdir = sys.argv[1] if len(sys.argv) > 1 else os.path.join(here, "gen1-pos")
    m = re.search(r"gen(\d+)-pos$", posdir.rstrip("/"))
    gen = m.group(1) if m else "1"
    outfile = sys.argv[2] if len(sys.argv) > 2 else os.path.join(here, "gen%s.html" % gen)
    sys.exit(build(posdir, outfile, gen))