"""Render the BHE_991 three-approach comparison contact sheet from saved runs.

Reads results/bhe_repeat.json (no model calls) and draws, per page:
GT (green), raw-vision C (red), data-grid A (orange), anatomy B-r36 (blue).
"""

from __future__ import annotations

import json
from pathlib import Path

from PIL import Image

from experiments.que245.gt import BHE_991_BANDS, ground_truth
from experiments.que245.overlay import draw_boxes, render_page

PDF = Path("documents/BHE_991.pdf")
RES = Path("experiments/que245/results/bhe_repeat.json")
OUT = Path("experiments/que245/render/final_compare_BHE.png")


def main():
    data = json.loads(RES.read_text())
    gt = ground_truth(PDF, BHE_991_BANDS)

    def run0(label):
        runs = data[label]["runs"]
        return {int(p): boxes[0] for p, boxes in runs.items()}

    c = run0("C raw-vision   ")
    a = run0("A data-grid r24")
    b = run0("B anatomy  r36")

    imgs = []
    for p in range(1, 9):
        im = render_page(PDF, p, dpi=120)
        im = draw_boxes(
            im,
            [
                ("GT", gt[p], (0, 160, 0)),
                ("C", c.get(p, []), (220, 0, 0)),
                ("A", a.get(p, []), (235, 140, 0)),
                ("B", b.get(p, []), (0, 80, 230)),
            ],
        )
        imgs.append(im)
    w = max(i.width for i in imgs)
    h = max(i.height for i in imgs)
    sheet = Image.new("RGB", (w * 4, h * 2), (255, 255, 255))
    for i, im in enumerate(imgs):
        sheet.paste(im, ((i % 4) * w, (i // 4) * h))
    sheet.save(OUT)
    print(f"saved {OUT} ({sheet.size})")
    print("legend: GT=green  C raw=red  A data-grid=orange  B anatomy=blue")


if __name__ == "__main__":
    main()