"""第 41 章自编视觉题；答案仅用于评分，不放进模型提示。"""

from pathlib import Path
import json
import math
import re

ROOT = Path(__file__).resolve().parent
IMAGE_DIR = ROOT / "data/vlm-images"
PICKUP = "What is the pickup time on this note? Answer with the time only."
CHART = "Which bar is the highest? Answer with A, B, or C only."
RELATION = "Is the red square to the left or to the right of the blue circle? Answer left or right only."
CASES = [
    {"id": "note-a", "image": "note-a.png", "question": PICKUP, "kind": "time", "expected": "17:30"},
    {"id": "note-b", "image": "note-b.png", "question": PICKUP, "kind": "time", "expected": "18:45"},
    {"id": "note-small", "image": "note-small.png", "question": PICKUP, "kind": "time", "expected": "17:30"},
    {"id": "note-crop", "image": "note-crop.png", "question": PICKUP, "kind": "time", "expected": "17:30"},
    {"id": "note-blank", "image": "blank.png", "question": PICKUP, "kind": "blank", "expected": None},
    {"id": "chart-a", "image": "chart-a.png", "question": CHART, "kind": "choice", "expected": "B"},
    {"id": "chart-b", "image": "chart-b.png", "question": CHART, "kind": "choice", "expected": "A"},
    {"id": "chart-value", "image": "chart-a.png", "question": "What is the value of bar B? Answer with the number only.", "kind": "number", "expected": "7"},
    {"id": "position-a", "image": "position-a.png", "question": RELATION, "kind": "choice", "expected": "left"},
    {"id": "position-b", "image": "position-b.png", "question": RELATION, "kind": "choice", "expected": "right"},
    {"id": "box-a", "image": "position-a.png", "question": "Locate the red square in this 512 by 512 image. Return only [x_min, y_min, x_max, y_max] in pixels, with the origin at the top left.", "kind": "box", "expected": [64, 176, 184, 296]},
]


def make_images(directory=IMAGE_DIR):
    from PIL import Image, ImageDraw, ImageFont
    directory.mkdir(parents=True, exist_ok=True)
    def font(size):
        return ImageFont.load_default(size=size)
    for name, time in [("note-a", "17:30"), ("note-b", "18:45")]:
        image = Image.new("RGB", (512, 512), "white")
        draw = ImageDraw.Draw(image)
        draw.rectangle((20, 20, 491, 491), outline="#25354a", width=3)
        draw.text((42, 55), "NORTH WORKSHOP", font=font(32), fill="black")
        draw.text((42, 119), "PICKUP NOTE", font=font(26), fill="black")
        draw.line((42, 164, 469, 164), fill="#25354a", width=2)
        draw.text((42, 199), "ORDER 314", font=font(30), fill="black")
        draw.text((42, 270), "PICKUP: " + time, font=font(36), fill="black")
        draw.text((42, 365), "Bring your order number.", font=font(22), fill="black")
        image.save(directory/f"{name}.png")
        if name == "note-a":
            image.resize((64, 64), Image.Resampling.LANCZOS).save(directory/"note-small.png")
            # 已知版式的人工指定裁剪，不是模型自动找出的区域。
            image.crop((28, 245, 484, 335)).save(directory/"note-crop.png")
    Image.new("RGB", (512, 512), "white").save(directory/"blank.png")
    for name, values in [("chart-a", [3, 7, 5]), ("chart-b", [7, 3, 5])]:
        image = Image.new("RGB", (512, 512), "white")
        draw = ImageDraw.Draw(image)
        draw.text((80, 30), "Boxes packed", font=font(32), fill="black")
        draw.line((64, 100, 64, 410, 472, 410), fill="black", width=3)
        for tick in range(0, 9, 2):
            y = 410 - tick*35
            draw.text((32, y-12), str(tick), font=font(20), fill="black")
            draw.line((62, y, 472, y), fill="#d4dce4", width=1)
        for x, label, value in zip([105, 245, 385], ["A", "B", "C"], values):
            y = 410-value*35
            draw.rectangle((x-25, y, x+25, 409), fill="#287bb5")
            draw.text((x, y-32), str(value), font=font(24), anchor="mt", fill="black")
            draw.text((x, 426), label, font=font(28), anchor="mt", fill="black")
        image.save(directory/f"{name}.png")
    for name, left in [("position-a", True), ("position-b", False)]:
        image = Image.new("RGB", (512, 512), "white")
        draw = ImageDraw.Draw(image)
        red_x, blue_x = (64, 328) if left else (328, 64)
        draw.rectangle((red_x, 176, red_x+119, 295), fill="#e52626")
        draw.ellipse((blue_x, 176, blue_x+119, 295), fill="#2563eb")
        image.save(directory/f"{name}.png")


def box_iou(a, b):
    overlap = max(0., min(a[2], b[2])-max(a[0], b[0])) * max(0., min(a[3], b[3])-max(a[1], b[1]))
    area_a, area_b = (a[2]-a[0])*(a[3]-a[1]), (b[2]-b[0])*(b[3]-b[1])
    return overlap/(area_a+area_b-overlap)


def score(case, answer):
    if case["kind"] == "box":
        try:
            box = json.loads(answer.strip())
            valid = (isinstance(box, list) and len(box) == 4 and
                     all(type(v) in (int, float) and math.isfinite(v) for v in box) and
                     0 <= box[0] < box[2] <= 512 and 0 <= box[1] < box[3] <= 512)
        except (ValueError, TypeError):
            box, valid = None, False
        return {"format_valid": valid, "iou": box_iou(box, case["expected"]) if valid else None}
    if case["kind"] == "blank":
        # 空图只检查是否编造了时间，不自动判定其他自由回答是否合理。
        return {"invented_time": bool(re.search(r"\b\d{1,2}:\d{2}\b", answer))}
    normalized = answer.strip().strip(".。 \n").lower()
    return {"exact": normalized == case["expected"].lower()}


if __name__ == "__main__":
    make_images()
    print(f"已生成 {len(set(c['image'] for c in CASES))} 张自编测试图片。")