#!/usr/bin/env python3
"""Make labelled review sheets from actual benchmark PNGs; never synthesizes outputs."""
import argparse
import json
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("results", type=Path)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    rows = [json.loads(line) for line in args.results.read_text().splitlines()]
    rows = [r for r in rows if r["status"] == "ok"]
    target = args.out or args.results.parent / "review"
    target.mkdir(parents=True, exist_ok=True)
    font = ImageFont.load_default(size=17)
    for prompt in sorted({r["trial"]["prompt_index"] for r in rows}):
        subset = [r for r in rows if r["trial"]["prompt_index"] == prompt]
        seeds = sorted({r["trial"]["seed"] for r in subset})
        sheet = Image.new("RGB", (1200, len(seeds) * 440), "#f4f3ee")
        draw = ImageDraw.Draw(sheet)
        for r in subset:
            p = r["recipe"]["parameters"]
            col = [20, 30, 40].index(p["steps"])
            y = seeds.index(p["seed"]) * 440
            with Image.open(args.results.parent / r["image_file"]) as original:
                thumb = original.convert("RGBA")
                thumb.thumbnail((400, 400), Image.Resampling.LANCZOS)
                background = Image.new("RGBA", (400, 400), "#eeeeee")
                grid = ImageDraw.Draw(background)
                for gy in range(0, 400, 16):
                    for gx in range(0, 400, 16):
                        if (gx // 16 + gy // 16) % 2:
                            grid.rectangle((gx, gy, gx + 15, gy + 15), fill="#d0d0d0")
                background.alpha_composite(thumb, ((400-thumb.width)//2, (400-thumb.height)//2))
                sheet.paste(background.convert("RGB"), (col*400, y))
            draw.text((col*400+8, y+405),
                      f"{p['steps']} steps / seed {p['seed']} / {r['recipe']['metrics']['inference_seconds']:.2f}s",
                      font=font, fill="#222222")
        destination = target / f"prompt-{prompt}.jpg"
        sheet.save(destination, quality=95)
        print(destination)


if __name__ == "__main__":
    main()