#!/usr/bin/env python3
"""
Restore missing ebl crops under hitit_ocr/data/classification/all/<label>/ebl_<image>_<idx>.png.

Legacy prepare_data.py load_ebl_coco parses COCO annotations and writes crops with:
  tablet_key = f"ebl_{img.file_name without ext}"
  crop_name  = f"{tablet_key}_{idx}.png"
where idx = position in anns_by_img[img_id].

We mirror that: parse the same two COCO jsons (train2017, val2017), build per-image annotation
lists in the same order as the legacy code, then for each manifest record whose path matches
ebl_<key>_<idx>.png, recompute the bbox, pad, and save.
"""
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path
from PIL import Image, ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

MIN_CROP_SIZE = 8
CONTEXT_PAD_RATIO = 0.15

PROJECT = Path("/arf/scratch/stakan/hitit-proje")
COCO_BASE = PROJECT / "datasets/sources/ebl_ocr/ready-for-training/coco-recognition/data/coco"
MANIFEST = PROJECT / "datasets/sources/hitit_local/manifest_v13_ultimate.jsonl"


def normalize_sign(s):
    if s is None:
        return None
    s = s.strip().rstrip(".,;: ").strip()
    if not s or s in ["/", ".", ",", "-", "(X)", "(x)", "X", "x", "?", ""]:
        return None
    if s.isdigit():
        return None
    if len(s) == 1 and not s.isalpha():
        return None
    return s


def build_coco_tablets():
    """Mirror load_ebl_coco return: {tablet_key: {image_path, width, height, signs[]}}."""
    tablets = {}
    for split, img_folder in [("instances_train2017.json", "train2017"),
                              ("instances_val2017.json", "val2017")]:
        ann_file = COCO_BASE / "annotations" / split
        img_dir = COCO_BASE / img_folder
        if not ann_file.exists():
            continue
        with open(ann_file) as f:
            coco = json.load(f)
        cat_map = {c["id"]: c["name"] for c in coco["categories"]}
        img_map = {img["id"]: img for img in coco["images"]}
        anns_by_img = defaultdict(list)
        for ann in coco["annotations"]:
            anns_by_img[ann["image_id"]].append(ann)
        for img_id, img_info in img_map.items():
            img_file = img_dir / img_info["file_name"]
            if not img_file.exists():
                continue
            signs = []
            for ann in anns_by_img.get(img_id, []):
                bx, by, bw, bh = ann["bbox"]
                cat_name = cat_map.get(ann["category_id"], "")
                cat_name = normalize_sign(cat_name) if cat_name else None
                if not cat_name:
                    continue
                if bw > 0 and bh > 0:
                    signs.append((bx, by, bw, bh, cat_name))
            if not signs:
                continue
            tablet_key = f"ebl_{img_info['file_name'].replace('.jpg', '').replace('.png', '')}"
            tablets[tablet_key] = {
                "image_path": str(img_file),
                "width": img_info["width"],
                "height": img_info["height"],
                "signs": signs,
            }
    return tablets


def missing_ebl_records():
    out = []
    with open(MANIFEST) as f:
        for line in f:
            r = json.loads(line)
            if r.get("source") != "ebl":
                continue
            p = r.get("path", "")
            if "/classification/all/" not in p:
                continue
            if Path(p).exists():
                continue
            out.append(r)
    return out


def parse_key_idx(stem):
    """'ebl_CBS.1515-0_165' -> ('ebl_CBS.1515-0', 165)."""
    parts = stem.rsplit("_", 1)
    if len(parts) != 2 or not parts[1].isdigit():
        return None, None
    return parts[0], int(parts[1])


def main():
    print("[plan] parsing COCO…")
    tablets = build_coco_tablets()
    print(f"  COCO tablets: {len(tablets)}")

    print("[plan] scanning manifest for missing ebl records…")
    missing = missing_ebl_records()
    print(f"  missing: {len(missing)}")

    # Group by tablet_key
    by_key = defaultdict(list)
    for r in missing:
        key, idx = parse_key_idx(Path(r["path"]).stem)
        if key is None:
            continue
        by_key[key].append((idx, r))

    stats = Counter()
    for key, items in by_key.items():
        info = tablets.get(key)
        if info is None:
            stats["no_tablet"] += len(items)
            continue
        img_w, img_h, signs = info["width"], info["height"], info["signs"]
        try:
            img = Image.open(info["image_path"]).convert("RGB")
        except Exception:
            stats["img_fail"] += len(items)
            continue
        for idx, r in items:
            if idx >= len(signs):
                stats["idx_oor"] += 1
                continue
            x, y, bw, bh, _ = signs[idx]
            px = bw * CONTEXT_PAD_RATIO
            py = bh * CONTEXT_PAD_RATIO
            x1 = max(0, int(x - px))
            y1 = max(0, int(y - py))
            x2 = min(img_w, int(x + bw + px))
            y2 = min(img_h, int(y + bh + py))
            if x2 - x1 < MIN_CROP_SIZE or y2 - y1 < MIN_CROP_SIZE:
                stats["too_small"] += 1
                continue
            out = Path(r["path"])
            out.parent.mkdir(parents=True, exist_ok=True)
            try:
                img.crop((x1, y1, x2, y2)).save(str(out))
                stats["written"] += 1
                man_w, man_h = r.get("width"), r.get("height")
                if man_w != x2 - x1 or man_h != y2 - y1:
                    stats["size_mismatch"] += 1
            except Exception:
                stats["save_fail"] += 1
        img.close()
    print("DONE:", dict(stats))


if __name__ == "__main__":
    main()