#!/usr/bin/env python3
"""Run the Tibetan modern-book layout detector (DocLayout-YOLO fine-tune) on
one or more page images, applying the recommended *per-class* confidence
thresholds.

The model is a 4-class DocLayout-YOLO (header, text-area, footnote, footer),
fine-tuned on the same `tam2col` labels as BDRC's primary RT-DETR-l release
(see the model card / blog post). Class ids are passed straight through
(this checkpoint was fine-tuned directly on our 4-class schema, unlike the
off-the-shelf DocStructBench checkpoint, which needs a class remap).

Usage:
    python infer.py --weights doclayout_yolo_tibetan_book_layout.pt --source page.jpg
    python infer.py --weights doclayout_yolo_tibetan_book_layout.pt --source pages/ --out preds
"""
from __future__ import annotations

import argparse
from pathlib import Path

IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}

# Per-class max-F1 operating points (see model card). Use --global-conf 0.30
# instead if you prefer one number for all classes.
CLASS_THRESHOLDS = {0: 0.42, 1: 0.66, 2: 0.27, 3: 0.48}
CONF_FLOOR = min(CLASS_THRESHOLDS.values())
NAMES = {0: "header", 1: "text-area", 2: "footnote", 3: "footer"}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--weights", required=True, help="path to the .pt weights")
    ap.add_argument("--source", required=True, help="image file or folder")
    ap.add_argument("--out", default=None,
                    help="optional folder to write YOLO-format .txt labels")
    ap.add_argument("--imgsz", type=int, default=1024)
    ap.add_argument("--device", default="0")
    ap.add_argument("--global-conf", type=float, default=None,
                    help="use ONE threshold for all classes instead of per-class")
    args = ap.parse_args()

    from doclayout_yolo import YOLOv10

    thresholds = ({c: args.global_conf for c in CLASS_THRESHOLDS}
                  if args.global_conf is not None else CLASS_THRESHOLDS)
    floor = min(thresholds.values())

    model = YOLOv10(args.weights)

    out_dir = Path(args.out) if args.out else None
    if out_dir:
        out_dir.mkdir(parents=True, exist_ok=True)

    results = model.predict(source=args.source, imgsz=args.imgsz, conf=floor,
                            device=args.device, stream=True, verbose=False)

    n_img = n_kept = 0
    for r in results:
        n_img += 1
        stem = Path(r.path).stem
        lines = []
        if r.boxes is not None:
            H, W = r.orig_shape
            for b, cf, cl in zip(r.boxes.xyxy.tolist(), r.boxes.conf.tolist(),
                                 r.boxes.cls.tolist()):
                cls = int(cl)
                if cls not in NAMES or cf < thresholds.get(cls, floor):
                    continue
                x1, y1, x2, y2 = b
                cx, cy = ((x1 + x2) / 2) / W, ((y1 + y2) / 2) / H
                w, h = (x2 - x1) / W, (y2 - y1) / H
                lines.append((cls, cf, cx, cy, w, h))
        n_kept += len(lines)
        print(f"{stem}: {len(lines)} boxes")
        for cls, cf, cx, cy, w, h in lines:
            print(f"    {NAMES[cls]:10} conf={cf:.3f}  "
                  f"cx={cx:.3f} cy={cy:.3f} w={w:.3f} h={h:.3f}")
        if out_dir:
            (out_dir / f"{stem}.txt").write_text(
                "".join(f"{c} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n"
                        for c, _, cx, cy, w, h in lines))

    print(f"\n{n_img} images, {n_kept} boxes kept (thresholds: {thresholds})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())