"""Blind pairwise comparison sheets for visual answers.

make:   blind_pairs.py make --a results/vis_base_dm/png --b results/vis_rep/png --ids v01 ... --out results/blind_rep
        Each row shows one task: left / right are A and B in a random order (seeded); the key is written to
        <out>/key.json and must not be opened before judging. Sheets: <out>/sheet_<k>.png, 8 tasks each.
score:  blind_pairs.py score --out results/blind_rep
        Reads <out>/judgments.json ({"v01": "L" | "R" | "=" , ...}, written before opening the key) and reports how
        often B (the candidate system) won, lost or tied against A.
"""
import argparse
import json
import os
import random
import re

from PIL import Image, ImageDraw

ap = argparse.ArgumentParser()
ap.add_argument("cmd", choices=["make", "score"])
ap.add_argument("--a")
ap.add_argument("--b")
ap.add_argument("--ids", nargs="+")
ap.add_argument("--out", required=True)
ap.add_argument("--sample", type=int, default=0)
ap.add_argument("--seed", type=int, default=7)
ap.add_argument("--w", type=int, default=420)
ap.add_argument("--per-sheet", type=int, default=8)
ap.add_argument("--tasks", help="task file: the first words of each task are printed over its row")
a = ap.parse_args()

if a.cmd == "make":
    heads = {}
    if a.tasks:
        for l in open(a.tasks, encoding="utf-8"):
            t = json.loads(l)
            heads[t["id"].split("/")[-1]] = t["prompt"][:70]
    os.makedirs(a.out, exist_ok=True)
    rng = random.Random(a.seed)
    key = {}
    w, h, pad = a.w, a.w * 3 // 4, 22
    rows = [a.ids[i:i + a.per_sheet] for i in range(0, len(a.ids), a.per_sheet)]
    for k, ids in enumerate(rows):
        sheet = Image.new("RGB", (2 * (w + 6) + 6, len(ids) * (h + pad + 6) + 6), "#202124")
        d = ImageDraw.Draw(sheet)
        for r, tid in enumerate(ids):
            flip = rng.random() < 0.5
            key[tid] = "B_left" if flip else "A_left"
            sides = [a.b, a.a] if flip else [a.a, a.b]
            for c, folder in enumerate(sides):
                x, y = 6 + c * (w + 6), 6 + r * (h + pad + 6)
                stem = tid if re.search(r"_\d+$", tid) else f"{tid}_{a.sample}"  # ids may name the sample: v01_1
                p = os.path.join(folder, f"{stem}.png")
                head = heads.get(re.sub(r"_\d+$", "", tid), "") if c == 0 else ""
                d.text((x + 4, y + 4), f"{tid}  {'L' if c == 0 else 'R'}  {head}", fill="#e8eaed")
                if os.path.exists(p):
                    sheet.paste(Image.open(p).convert("RGB").resize((w, h)), (x, y + pad))
                else:
                    d.rectangle([x, y + pad, x + w, y + pad + h], outline="#5f6368")
                    d.text((x + w // 2 - 30, y + pad + h // 2), "no output", fill="#9aa0a6")
        sheet.save(os.path.join(a.out, f"sheet_{k}.png"))
    json.dump(key, open(os.path.join(a.out, "key.json"), "w"))
    print(f"{len(rows)} sheets in {a.out}; key written (do not open before judging)")
else:
    key = json.load(open(os.path.join(a.out, "key.json")))
    judg = json.load(open(os.path.join(a.out, "judgments.json")))
    win = lose = tie = 0
    for tid, j in judg.items():
        if j == "=":
            tie += 1
            continue
        b_left = key[tid] == "B_left"
        if (j == "L") == b_left:
            win += 1
        else:
            lose += 1
    n = win + lose + tie
    print(f"candidate B vs A over {n} tasks: wins {win}, losses {lose}, ties {tie}  "
          f"(win rate excluding ties {100 * win / max(1, win + lose):.0f}%)")