#!/usr/bin/env python3
"""CPU audit of the released Figure 3 heatmap artifacts.

The paper archive contains the two rendered 32B heatmaps but no checkpoint.
This audit extracts the figures, renders them at a fixed resolution, samples
all 20 x 10 cells, and calibrates each cell against the figure's own 0--1
colorbar.  It is deliberately an artifact audit, not a substituted model run.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import tarfile
import tempfile
from pathlib import Path


FIGURES = {
    "baseline": "figs/needle_haystack_heatmap_chinese_32b_baseline.pdf",
    "ours": "figs/needle_haystack_heatmap_chinese_32b_ours.pdf",
}
WIDTH, HEIGHT = 1666, 870
GRID_X = [185 + 64 * i for i in range(20)]
GRID_Y = [62 + 70 * i for i in range(10)]
BAR_X = 1530
BAR_Y = list(range(27, 728))


def read_ppm(path: Path) -> tuple[int, int, bytes]:
    raw = path.read_bytes()
    tokens = []
    at = 0
    while len(tokens) < 4:
        while raw[at] in b" \t\r\n":
            at += 1
        end = at
        while raw[end] not in b" \t\r\n":
            end += 1
        tokens.append(raw[at:end])
        at = end
    magic, width, height, maxval = tokens
    if magic != b"P6" or maxval != b"255":
        raise ValueError("expected an 8-bit binary PPM")
    while raw[at] in b" \t\r\n":
        at += 1
    return int(width), int(height), raw[at:]


def rgb(data: bytes, width: int, x: int, y: int) -> tuple[int, int, int]:
    at = (y * width + x) * 3
    return tuple(data[at : at + 3])


def extract(tar: tarfile.TarFile, name: str, destination: Path) -> None:
    for member in tar:
        if member.name.endswith(name):
            handle = tar.extractfile(member)
            if handle is None:
                break
            destination.write_bytes(handle.read())
            return
    raise FileNotFoundError(name)


def audit_figure(pdf: Path, scratch: Path) -> dict:
    prefix = scratch / pdf.stem
    subprocess.run(
        ["pdftoppm", "-r", "150", "-singlefile", str(pdf), str(prefix)],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    width, height, data = read_ppm(prefix.with_suffix(".ppm"))
    if (width, height) != (WIDTH, HEIGHT):
        raise AssertionError((width, height))
    bar = [rgb(data, width, BAR_X, y) for y in BAR_Y]
    scores = []
    color_rows = []
    for y in GRID_Y:
        row = []
        for x in GRID_X:
            observed = rgb(data, width, x, y)
            nearest = min(
                range(len(bar)),
                key=lambda j: sum((observed[k] - bar[j][k]) ** 2 for k in range(3)),
            )
            score = round(1.0 - nearest / (len(bar) - 1), 1)
            row.append(score)
            scores.append(score)
        color_rows.append(row)
    return {
        "cells": len(scores),
        "mean": sum(scores) / len(scores),
        "mean_percent": 100.0 * sum(scores) / len(scores),
        "histogram": {str(value): scores.count(value) for value in sorted(set(scores))},
        "rows": color_rows,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--archive", default="source/2603.21719v1.tar")
    args = parser.parse_args()
    with tempfile.TemporaryDirectory(prefix="tablelong-heatmap-") as directory:
        scratch = Path(directory)
        results = {}
        with tarfile.open(args.archive, "r:*") as tar:
            for label, name in FIGURES.items():
                pdf = scratch / f"{label}.pdf"
                extract(tar, name, pdf)
                results[label] = audit_figure(pdf, scratch)
    if results["baseline"]["mean_percent"] != 87.95:
        raise AssertionError(results)
    if results["ours"]["mean_percent"] != 99.4:
        raise AssertionError(results)
    print(json.dumps(results, sort_keys=True))


if __name__ == "__main__":
    main()