histogram.py

"""Static PNG renders of a corpus.

Two forms:
  • transcript — sequential, every token rendered as a colored block, verse
    newlines respected. The corpus printed in pure color, like a manuscript page.
  • palette — frequency-weighted bars sorted by hue. The corpus's chromatic
    fingerprint: a Pantone fan extracted from the text.
"""

import colorsys
import math
from collections import Counter
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

from engine import load_lut, load_tokenizer


# ---------- transcript ---------- #

def _build_cube(text: str, cell_size: int, gap: int, bg) -> Image.Image:
    """Bare cube of color cells, no outer margin. side × side perfect square."""
    tokenizer = load_tokenizer()
    lut, _ = load_lut()
    ids = tokenizer.encode(text, add_special_tokens=False)
    tokens = [tid for tid in ids if tokenizer.decode([tid]).strip()]
    n = len(tokens)
    if n == 0:
        return Image.new("RGB", (cell_size, cell_size), bg)

    side = math.ceil(math.sqrt(n))
    w = side * cell_size + max(0, side - 1) * gap
    img = Image.new("RGB", (w, w), bg)
    draw = ImageDraw.Draw(img)
    for i, tid in enumerate(tokens):
        row, col = divmod(i, side)
        x = col * (cell_size + gap)
        y = row * (cell_size + gap)
        r, g, b = lut[tid]
        draw.rectangle(
            [x, y, x + cell_size - 1, y + cell_size - 1],
            fill=(int(r), int(g), int(b)),
        )
    return img


def _build_palette_strip(text: str, width: int, height: int, bg) -> Image.Image:
    """Frequency-weighted color bars sorted by hue. Width should match the cube's."""
    tokenizer = load_tokenizer()
    lut, _ = load_lut()
    ids = tokenizer.encode(text, add_special_tokens=False)
    counts = Counter(ids)
    items = [(tid, c) for tid, c in counts.items()]
    items.sort(key=lambda x: _hue_key(lut[x[0]]))
    total = sum(c for _, c in items) or 1

    img = Image.new("RGB", (width, height), bg)
    draw = ImageDraw.Draw(img)
    cursor = 0.0
    for tid, c in items:
        bar_w = (c / total) * width
        if bar_w < 0.5:
            cursor += bar_w
            continue
        x0 = int(round(cursor))
        x1 = int(round(cursor + bar_w))
        if x1 <= x0:
            x1 = x0 + 1
        r, g, b = lut[tid]
        draw.rectangle([x0, 0, x1 - 1, height - 1], fill=(int(r), int(g), int(b)))
        cursor += bar_w
    return img


def render_transcript(
    text: str,
    cell_size: int = 18,
    gap: int = 2,
    palette_h: int = 52,
    pad_top: int = 26,
    pad_sides: int = 26,
    pad_between: int = 18,
    pad_bottom: int = 64,
    frame=(244, 240, 232),     # warm cream — polaroid paper
    cube_bg=(8, 8, 8),
) -> Image.Image:
    """Polaroid-style render: square color cube on top, hue-sorted palette strip
    below, cream frame around the whole thing (thin on top/sides, thick on the
    bottom — like a polaroid where you'd write the caption)."""
    cube = _build_cube(text, cell_size, gap, cube_bg)
    palette = _build_palette_strip(text, cube.width, palette_h, cube_bg)

    inner_w = cube.width
    inner_h = cube.height + pad_between + palette.height
    img_w = inner_w + 2 * pad_sides
    img_h = pad_top + inner_h + pad_bottom

    img = Image.new("RGB", (img_w, img_h), frame)
    img.paste(cube, (pad_sides, pad_top))
    img.paste(palette, (pad_sides, pad_top + cube.height + pad_between))
    return img


# ---------- palette / histogram ---------- #

def _hue_key(rgb):
    r, g, b = (c / 255 for c in rgb)
    h, s, v = colorsys.rgb_to_hsv(r, g, b)
    # group by hue band, then by value/saturation within band
    return (round(h * 24), -v, -s)


def render_palette(
    text: str,
    width: int = 1400,
    height: int = 220,
    bg=(8, 8, 8),
) -> Image.Image:
    """Standalone palette: full-width chromatic ribbon, the corpus's color signature."""
    return _build_palette_strip(text, width, height, bg)


_FONT_CANDIDATES = [
    "consola.ttf", "Consolas.ttf", "consolab.ttf",
    "DejaVuSansMono.ttf", "DejaVuSansMono-Bold.ttf",
    "cour.ttf", "Courier New.ttf",
]
_FONT_CACHE: dict = {}


def _try_font(name_or_path: str, size: int):
    key = (name_or_path, size)
    if key in _FONT_CACHE:
        return _FONT_CACHE[key]
    try:
        font = ImageFont.truetype(name_or_path, size)
    except (OSError, IOError):
        font = None
    _FONT_CACHE[key] = font
    return font


def _resolve_font_path() -> str | None:
    """Walk the candidate list once and return the first that loads."""
    for name in _FONT_CANDIDATES:
        if _try_font(name, 12) is not None:
            return name
    return None


def _fit_text(draw, text: str, font_path: str, max_w: int, max_h: int,
              min_size: int = 5, max_size: int = 240):
    """Pick the largest font size where `text` fits the box. Truncates with '…'
    if even at min_size it overflows. Returns (font, displayed_text)."""
    if not font_path:
        return ImageFont.load_default(), text[:1] or "·"

    # Step coarsely (4pt jumps) then refine — much faster than dense scan
    fits_size = None
    for size in range(max_size, min_size - 1, -4):
        font = _try_font(font_path, size)
        if font is None:
            continue
        bb = draw.textbbox((0, 0), text, font=font)
        if bb[2] - bb[0] <= max_w and bb[3] - bb[1] <= max_h:
            fits_size = size
            break

    if fits_size is not None:
        # Refine upward by 1pt steps to squeeze the last bit
        for size in range(fits_size + 1, min(fits_size + 4, max_size) + 1):
            font = _try_font(font_path, size)
            if font is None:
                continue
            bb = draw.textbbox((0, 0), text, font=font)
            if bb[2] - bb[0] <= max_w and bb[3] - bb[1] <= max_h:
                fits_size = size
        return _try_font(font_path, fits_size), text

    # Doesn't fit at any size — truncate at min_size
    font = _try_font(font_path, min_size) or ImageFont.load_default()
    for cut in range(len(text) - 1, 0, -1):
        candidate = text[:cut] + "…"
        bb = draw.textbbox((0, 0), candidate, font=font)
        if bb[2] - bb[0] <= max_w and bb[3] - bb[1] <= max_h:
            return font, candidate
    return font, "…"


def _display_token(tok_text: str) -> str:
    """Make whitespace visible: leading space → '·', newline → '↵', tab → '→'."""
    s = tok_text.replace("\n", "↵").replace("\t", "→")
    if s.startswith(" "):
        s = "·" + s[1:]
    return s


def render_reading_map(
    text: str,
    title: str = "",
    cell_mm: float | None = None,
    dpi: int = 300,
    margin_mm: float = 12.0,
    font_path: str | None = None,
    bg=(244, 240, 232),
) -> Image.Image:
    """Per-corpus reading-along map — colored grid mirroring the playground's
    per-corpus layout, with each cell labeled with its token text. Auto-picks
    cell size: 25mm if ≤500 unique tokens (handout), 15mm otherwise (wall).
    Output is print-ready at the chosen DPI.
    """
    tokenizer = load_tokenizer()
    lut, _ = load_lut()

    ids = tokenizer.encode(text, add_special_tokens=False)
    unique = sorted({tid for tid in ids if tokenizer.decode([tid]).strip()})
    n = len(unique)

    if cell_mm is None:
        cell_mm = 25.0 if n <= 500 else 15.0

    def mm(v: float) -> int:
        return max(1, round(v * dpi / 25.4))

    cell_px = mm(cell_mm)
    margin_px = mm(margin_mm)

    if n == 0:
        empty = Image.new("RGB", (margin_px * 2, margin_px * 2), bg)
        return empty

    side = math.ceil(math.sqrt(n))
    grid_px = side * cell_px

    title_band = mm(margin_mm * 1.4)
    img_w = grid_px + 2 * margin_px
    img_h = grid_px + margin_px + title_band

    img = Image.new("RGB", (img_w, img_h), bg)
    draw = ImageDraw.Draw(img)

    if font_path is None:
        font_path = _resolve_font_path()

    pad = max(2, cell_px // 14)
    inner_w = cell_px - 2 * pad
    inner_h = cell_px - 2 * pad

    grid_top = title_band
    grid_left = margin_px

    for i, tid in enumerate(unique):
        row, col = divmod(i, side)
        x0 = grid_left + col * cell_px
        y0 = grid_top + row * cell_px
        r, g, b = (int(c) for c in lut[tid])
        draw.rectangle([x0, y0, x0 + cell_px - 1, y0 + cell_px - 1], fill=(r, g, b))

        disp = _display_token(tokenizer.decode([tid]))
        if not disp:
            continue
        # Estimate sensible max font size: ~1.6 px per char-width is typical for
        # monospace, so this skips obviously-too-big sizes for long tokens
        max_estimate = min(inner_h, int(inner_w / max(1, len(disp)) * 1.7))
        font, shown = _fit_text(draw, disp, font_path, inner_w, inner_h,
                                min_size=5, max_size=max(8, max_estimate))
        bb = draw.textbbox((0, 0), shown, font=font)
        tw, th = bb[2] - bb[0], bb[3] - bb[1]
        tx = x0 + (cell_px - tw) // 2 - bb[0]
        ty = y0 + (cell_px - th) // 2 - bb[1]
        lum = 0.2126 * r + 0.7152 * g + 0.0722 * b
        fg = (0, 0, 0) if lum > 140 else (245, 240, 230)
        draw.text((tx, ty), shown, font=font, fill=fg)

    # Title strip
    title_size = max(10, cell_px // 4)
    title_font = _try_font(font_path, title_size) if font_path else None
    if title_font is None:
        title_font = ImageFont.load_default()
    label = title or ""
    suffix = f"{n} unique tokens  ·  cell {cell_mm:g} mm  ·  {dpi} dpi"
    full = f"{label}    {suffix}" if label else suffix
    bb = draw.textbbox((0, 0), full, font=title_font)
    tw, th = bb[2] - bb[0], bb[3] - bb[1]
    tx = (img_w - tw) // 2 - bb[0]
    ty = (title_band - th) // 2 - bb[1]
    draw.text((tx, ty), full, font=title_font, fill=(70, 60, 50))

    return img


def render_vocab_map(
    cell: int = 6,
    gap: int = 0,
    sort: str = "id",
    bg=(8, 8, 8),
) -> Image.Image:
    """The whole vocabulary as a chromatic atlas — every token rendered as one
    cell, packed into a √V × √V grid.

    `sort`:
      - "id"  — token ID order (roughly BPE merge frequency: common merges first,
                 then progressively rarer / more-CJK-heavy tokens). Reveals the
                 model's *training-derived* vocabulary structure.
      - "hue" — rainbow-ordered. Reveals the *chromatic* distribution: what
                 fraction of the model's vocabulary occupies each color region.
    """
    lut, meta = load_lut()
    vocab = meta["vocab_size"]

    if sort == "hue":
        order = sorted(range(vocab), key=lambda i: _hue_key(tuple(lut[i])))
    else:
        order = range(vocab)

    side = math.ceil(math.sqrt(vocab))
    img_w = side * cell + max(0, side - 1) * gap
    img = Image.new("RGB", (img_w, img_w), bg)
    draw = ImageDraw.Draw(img)

    for i, tid in enumerate(order):
        row, col = divmod(i, side)
        x = col * (cell + gap)
        y = row * (cell + gap)
        r, g, b = lut[tid]
        draw.rectangle(
            [x, y, x + cell - 1, y + cell - 1],
            fill=(int(r), int(g), int(b)),
        )
    return img