#!/usr/bin/env python3
"""Run the TiBLA PP-DocLayout-L Tibetan book layout detector on one or more page
images, writing YOLO-format labels.

The model is a 4-class PP-DocLayout-L (header, text-area, footnote, footer)
fine-tuned on the leak-free v4 `tam2col` split of TiBLAD. It is served from the
exported PaddlePaddle inference model in `inference/`. The recommended global
operating confidence is ~0.68 (the best-mean-F1 point on the v4 test).

Usage:
    python infer.py --model-dir inference --source page.jpg
    python infer.py --model-dir inference --source pages/ --out preds --conf 0.68

Requires: paddlepaddle==3.0.0, paddlex.
"""
from __future__ import annotations

import argparse
from pathlib import Path

IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
# native class order from inference.yml
NAMES = {"header": 0, "text-area": 1, "footnote": 2, "footer": 3}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model-dir", default="inference",
                    help="path to the exported inference model dir")
    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("--conf", type=float, default=0.68,
                    help="global operating confidence (default 0.68)")
    args = ap.parse_args()

    from PIL import Image
    from paddlex import create_model

    model = create_model(model_name="PP-DocLayout-L", model_dir=args.model_dir)

    src = Path(args.source)
    imgs = sorted(p for p in src.iterdir() if p.suffix.lower() in IMG_EXTS) \
        if src.is_dir() else [src]
    out_dir = Path(args.out) if args.out else None
    if out_dir:
        out_dir.mkdir(parents=True, exist_ok=True)

    n_img = n_kept = 0
    for ip in imgs:
        n_img += 1
        with Image.open(ip) as im:
            W, H = im.size
        lines = []
        for res in model.predict(str(ip), threshold=args.conf):
            for box in res["boxes"]:
                label = box["label"]
                cls = NAMES.get(label)
                if cls is None:
                    continue
                x1, y1, x2, y2 = box["coordinate"]
                cx, cy = ((x1 + x2) / 2) / W, ((y1 + y2) / 2) / H
                w, h = (x2 - x1) / W, (y2 - y1) / H
                lines.append((cls, float(box["score"]), cx, cy, w, h))
        n_kept += len(lines)
        print(f"{ip.stem}: {len(lines)} boxes")
        for cls, score, cx, cy, w, h in lines:
            name = next(k for k, v in NAMES.items() if v == cls)
            print(f"    {name:10} conf={score:.3f}  cx={cx:.3f} cy={cy:.3f} "
                  f"w={w:.3f} h={h:.3f}")
        if out_dir:
            (out_dir / f"{ip.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 (conf {args.conf})")
    return 0


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