benchmark_grid_draw.py

#!/usr/bin/env python3

from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, List, Sequence, Tuple

from PIL import Image, ImageDraw, ImageFont


COLUMNS = 200
CELL_WIDTH = 4  # inner fill width
CELL_HEIGHT = 8  # inner fill height
BORDER = 1  # divider thickness between cells
ROW_HEIGHT = CELL_HEIGHT + BORDER
GRID_WIDTH = BORDER + COLUMNS * (CELL_WIDTH + BORDER)
GRID_ORIGIN_Y = 0
LABEL_LEFT_PADDING = 12
LABEL_RIGHT_PADDING = 12
TITLE_TOP_PADDING = 6
TITLE_BOTTOM_PADDING = 8
GUIDE_TOP_PADDING = 6
GUIDE_LINE_SPACING = 4
GUIDE_BOTTOM_PADDING = 4
KEY_TOP_PADDING = 8
KEY_BOTTOM_PADDING = 8
KEY_BOX_WIDTH = 12
KEY_BOX_HEIGHT = 12
KEY_ITEM_SPACING = 18
KEY_BOX_TEXT_GAP = 6
COLORS = {
    0: (255, 255, 255),  # white
    1: (200, 200, 200),  # light grey
    2: (0, 128, 0),  # green
    3: (65, 105, 225),  # blue
    4: (160, 32, 240),  # purple
}
BORDER_COLOR = (180, 180, 180)
TEXT_COLOR = (0, 0, 0)
COLOR_KEY_ITEMS = [
    (0, "0 languages (unsolved)"),
    (1, "solved in 1 language"),
    (2, "2 languages"),
    (3, "3 languages"),
    (4, "4 languages"),
]


def load_benchmark_data(path: Path) -> Dict[str, Dict[str, str]]:
    raw = path.read_text()
    data = json.loads(raw)
    if not isinstance(data, dict):
        raise ValueError("Top-level JSON structure must be an object")
    return data


def collect_models_with_tests(
    data: Dict[str, Dict[str, object]]
) -> List[Tuple[str, List[str], bool]]:
    models: List[Tuple[str, List[str], bool]] = []
    required_keys = [f"{lang}-200-test" for lang in ("python", "java", "rust", "clojure")]

    for name, metrics in data.items():
        if not isinstance(metrics, dict):
            continue
        values: List[str] = []
        for key in required_keys:
            value = metrics.get(key)
            if not (isinstance(value, str) and len(value) == COLUMNS and set(value) <= {"0", "1"}):
                break
            values.append(value)
        else:
            thinking_value = metrics.get("thinking", False)
            if isinstance(thinking_value, str):
                is_thinking = thinking_value.strip().lower() == "true"
            else:
                is_thinking = bool(thinking_value)
            models.append((name, values, is_thinking))

    return models


def count_solutions(values: List[str]) -> List[int]:
    aggregated = []
    for idx in range(COLUMNS):
        count = sum(int(value[idx]) for value in values)
        aggregated.append(count)
    return aggregated


def color_for_count(count: int) -> Tuple[int, int, int]:
    if count >= 4:
        return COLORS[4]
    return COLORS[count]


def measure_text_width(font: ImageFont.ImageFont, text: str) -> int:
    try:
        bbox = font.getbbox(text)
    except AttributeError:
        width, _ = font.getsize(text)
        return width
    return bbox[2] - bbox[0]


def measure_text_height(font: ImageFont.ImageFont, text: str) -> int:
    try:
        bbox = font.getbbox(text)
    except AttributeError:
        _, height = font.getsize(text)
        return height
    return bbox[3] - bbox[1]


def load_font() -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    # Try to load a compact monospace font; fall back to default.
    try:
        return ImageFont.truetype("DejaVuSansMono.ttf", size=8)
    except OSError:
        return ImageFont.load_default()


def draw_grid(
    draw: ImageDraw.ImageDraw,
    rows: int,
    colors: List[List[Tuple[int, int, int]]],
    origin_x: int,
    origin_y: int,
) -> None:
    grid_width = GRID_WIDTH
    for row_index in range(rows):
        for col_index in range(COLUMNS):
            x0 = origin_x + BORDER + col_index * (CELL_WIDTH + BORDER)
            y0 = origin_y + BORDER + row_index * (CELL_HEIGHT + BORDER)
            x1 = x0 + CELL_WIDTH - 1
            y1 = y0 + CELL_HEIGHT - 1
            draw.rectangle((x0, y0, x1, y1), fill=colors[row_index][col_index])

    grid_height = BORDER + rows * (CELL_HEIGHT + BORDER)

    # Vertical grid lines
    for offset in range(0, grid_width, CELL_WIDTH + BORDER):
        x = origin_x + offset
        draw.line((x, origin_y, x, origin_y + grid_height - 1), fill=BORDER_COLOR)

    # Right-most border line
    draw.line(
        (
            origin_x + grid_width - 1,
            origin_y,
            origin_x + grid_width - 1,
            origin_y + grid_height - 1,
        ),
        fill=BORDER_COLOR,
    )

    # Horizontal grid lines
    for offset in range(0, grid_height, CELL_HEIGHT + BORDER):
        y = origin_y + offset
        draw.line((origin_x, y, origin_x + grid_width - 1, y), fill=BORDER_COLOR)

    # Bottom border line
    draw.line(
        (
            origin_x,
            origin_y + grid_height - 1,
            origin_x + grid_width - 1,
            origin_y + grid_height - 1,
        ),
        fill=BORDER_COLOR,
    )


def add_labels(
    draw: ImageDraw.ImageDraw,
    models: Sequence[str],
    font: ImageFont.ImageFont,
    left_padding: int,
    origin_y: int,
) -> None:
    for index, name in enumerate(models):
        row_center_y = origin_y + BORDER + CELL_HEIGHT / 2 + index * ROW_HEIGHT
        draw.text((left_padding, row_center_y), name, font=font, fill=TEXT_COLOR, anchor="lm")


def draw_column_guides(
    draw: ImageDraw.ImageDraw,
    font: ImageFont.ImageFont,
    label_width: int,
    grid_height: int,
    origin_y: int,
) -> None:
    text_height = measure_text_height(font, "0")
    line1_y = origin_y + grid_height + GUIDE_TOP_PADDING + text_height / 2
    line2_y = line1_y + text_height + GUIDE_LINE_SPACING

    for col in range(COLUMNS):
        center_x = label_width + BORDER + col * (CELL_WIDTH + BORDER) + CELL_WIDTH / 2
        digit_line1 = str((col + 1) % 10)
        draw.text((center_x, line1_y), digit_line1, font=font, fill=TEXT_COLOR, anchor="mm")

        if (col + 1) % 10 == 0:
            digit_line2 = str(((col + 1) // 10) % 10)
            draw.text((center_x, line2_y), digit_line2, font=font, fill=TEXT_COLOR, anchor="mm")


def draw_color_key(
    draw: ImageDraw.ImageDraw,
    font: ImageFont.ImageFont,
    origin_x: int,
    origin_y: int,
    max_width: int,
) -> int:
    x = origin_x
    y = origin_y
    row_height = max(KEY_BOX_HEIGHT, measure_text_height(font, "Ag"))
    center_y = y + row_height / 2

    for count, label in COLOR_KEY_ITEMS:
        label_width = measure_text_width(font, label)
        item_width = KEY_BOX_WIDTH + KEY_BOX_TEXT_GAP + label_width
        if x != origin_x and x + item_width > origin_x + max_width:
            x = origin_x
            y += row_height + GUIDE_LINE_SPACING
            center_y = y + row_height / 2

        box_top = center_y - KEY_BOX_HEIGHT / 2
        box_left = x
        draw.rectangle((box_left, box_top, box_left + KEY_BOX_WIDTH - 1, box_top + KEY_BOX_HEIGHT - 1), fill=color_for_count(count), outline=BORDER_COLOR)
        draw.text((box_left + KEY_BOX_WIDTH + KEY_BOX_TEXT_GAP, center_y), label, font=font, fill=TEXT_COLOR, anchor="lm")

        x += item_width + KEY_ITEM_SPACING

    return int(y + row_height - origin_y)


def measure_color_key_height(font: ImageFont.ImageFont, max_width: int) -> int:
    x = 0
    rows = 1
    row_height = max(KEY_BOX_HEIGHT, measure_text_height(font, "Ag"))

    for _, label in COLOR_KEY_ITEMS:
        label_width = measure_text_width(font, label)
        item_width = KEY_BOX_WIDTH + KEY_BOX_TEXT_GAP + label_width
        if x != 0 and x + item_width > max_width:
            rows += 1
            x = 0
        x += item_width + KEY_ITEM_SPACING

    return int(rows * row_height + (rows - 1) * GUIDE_LINE_SPACING)


def create_image(
    models_colors: List[List[Tuple[int, int, int]]],
    model_names: List[str],
    font: ImageFont.ImageFont,
    label_width: int,
    output_path: Path,
    title: str,
) -> None:
    rows = len(models_colors)
    title_height = measure_text_height(font, title)
    title_block_height = TITLE_TOP_PADDING + title_height + TITLE_BOTTOM_PADDING

    grid_height = BORDER + rows * (CELL_HEIGHT + BORDER)
    text_height = measure_text_height(font, "0")
    guide_height = GUIDE_TOP_PADDING + text_height * 2 + GUIDE_LINE_SPACING + GUIDE_BOTTOM_PADDING
    key_height = KEY_TOP_PADDING + measure_color_key_height(font, GRID_WIDTH) + KEY_BOTTOM_PADDING
    total_height = title_block_height + grid_height + guide_height + key_height

    img = Image.new("RGB", (label_width + GRID_WIDTH, int(total_height)), "white")
    draw = ImageDraw.Draw(img)

    total_width = label_width + GRID_WIDTH
    title_center_x = total_width / 2
    title_center_y = TITLE_TOP_PADDING + title_height / 2
    draw.text((title_center_x, title_center_y), title, font=font, fill=TEXT_COLOR, anchor="mm")

    grid_origin_y = title_block_height
    draw_grid(draw, rows, models_colors, origin_x=label_width, origin_y=grid_origin_y)
    add_labels(draw, model_names, font, left_padding=LABEL_LEFT_PADDING, origin_y=grid_origin_y)
    draw_column_guides(draw, font, label_width, grid_height, origin_y=grid_origin_y)
    draw_color_key(
        draw,
        font,
        origin_x=label_width,
        origin_y=grid_origin_y + grid_height + guide_height + KEY_TOP_PADDING,
        max_width=GRID_WIDTH,
    )

    img.save(output_path)


def main() -> None:
    benchmark_path = Path(__file__).with_name("benchmark.json")
    output_instruct_path = Path(__file__).with_name("benchmark_instruct.png")
    output_thinking_path = Path(__file__).with_name("benchmark_thinking.png")

    data = load_benchmark_data(benchmark_path)
    models = collect_models_with_tests(data)

    if not models:
        raise SystemExit("No models with complete 200-test data found.")

    font = load_font()
    max_label_width = max(measure_text_width(font, name) for name, _, _ in models)
    label_width = LABEL_LEFT_PADDING + max_label_width + LABEL_RIGHT_PADDING

    instruct = [(name, values) for name, values, is_thinking in models if not is_thinking]
    thinking = [(name, values) for name, values, is_thinking in models if is_thinking]

    def render_group(group: List[Tuple[str, List[str]]], output_path: Path, title: str) -> None:
        group_names = [name for name, _ in group]
        group_colors: List[List[Tuple[int, int, int]]] = []
        for _, values in group:
            counts = count_solutions(values)
            row_colors = [color_for_count(count) for count in counts]
            group_colors.append(row_colors)

        create_image(group_colors, group_names, font, label_width, output_path, title=title)
        print(f"Rendered benchmark grid for {len(group)} models to {output_path.name}")

    render_group(instruct, output_instruct_path, "Project Euler LLM Benchmark: instruct models")
    render_group(thinking, output_thinking_path, "Project Euler LLM Benchmark: thinking models")


if __name__ == "__main__":
    main()