"""Crop photographs to the zero-shot inscription box (data/ground_4b.jsonl) with a 20% margin and
rescale each crop to about 0.64 MP, so that runes occupy more of the recognizer's input.
Boxes are re-parsed from the raw model output (the "box" field of runs before the parse_box fix is wrong).
Images without a usable box are kept whole. Writes crops/ and splits/ocr_crop_{train,val,test}.jsonl.
Runs as the first step of the ocr4b_crop job (exp_v4.sh)."""
import json
import os
import re

from PIL import Image

R = "/home/jovyan/humonen/runic_access"
TARGET = 640_000
MARGIN = 0.2


def parse_box(text, w, h):  # same as ground.parse_box, without importing torch
    m = re.search(r"\[([^\[\]]*)\]", text)
    nums = re.findall(r"-?\d+(?:\.\d+)?", m.group(1) if m else "")
    if len(nums) < 4:
        return None
    x1, y1, x2, y2 = [float(v) for v in nums[:4]]
    sx, sy = (w, h) if max(x1, y1, x2, y2) <= 1.0 else (w / 1000.0, h / 1000.0)
    x1, x2 = sorted((x1 * sx, x2 * sx))
    y1, y2 = sorted((y1 * sy, y2 * sy))
    x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
    if x2 - x1 < 0.05 * w or y2 - y1 < 0.05 * h:
        return None
    return [x1, y1, x2, y2]


def crop_one(src, raw_text, dst):
    """Crop src to the parsed box with a margin, rescale to ~TARGET pixels, save to dst; return the crop box or None."""
    im = Image.open(src).convert("RGB")
    box = parse_box(raw_text, im.width, im.height)
    if box:
        x1, y1, x2, y2 = box
        mx, my = MARGIN * (x2 - x1), MARGIN * (y2 - y1)
        box = [max(0, x1 - mx), max(0, y1 - my), min(im.width, x2 + mx), min(im.height, y2 + my)]
        im = im.crop(tuple(round(v) for v in box))
    s = (TARGET / (im.width * im.height)) ** 0.5
    im = im.resize((max(28, round(im.width * s)), max(28, round(im.height * s))), Image.BICUBIC)
    im.save(dst + ".tmp.jpg", quality=92)
    os.replace(dst + ".tmp.jpg", dst)
    return box


if __name__ == "__main__":
    raw = {}
    for l in open(f"{R}/data/ground_4b.jsonl"):
        r = json.loads(l)
        raw[r["image"]] = r["raw"]
    os.makedirs(f"{R}/crops", exist_ok=True)
    stats = {"cropped": 0, "whole": 0}
    boxes_out = open(f"{R}/data/crop_boxes.jsonl", "w")
    for split in ("train", "val", "test"):
        with open(f"{R}/splits/ocr_crop_{split}.jsonl", "w") as f:
            for l in open(f"{R}/splits/ocr_real_{split}.jsonl"):
                r = json.loads(l)
                dst_rel = "crops/" + os.path.basename(r["image"])
                box = crop_one(f"{R}/{r['image']}", raw.get(r["image"], ""), f"{R}/{dst_rel}")
                stats["cropped" if box else "whole"] += 1
                boxes_out.write(json.dumps({"image": r["image"], "crop_box": box}) + "\n")
                f.write(json.dumps(dict(r, image=dst_rel, image_full=r["image"]), ensure_ascii=False) + "\n")
        print(split, stats, flush=True)
    boxes_out.close()
    print("DONE", stats)