# test_ocr_runtime.py
# 사용 예:
#  python test_ocr_runtime.py --imgs img --onnx_dir PP_Models/onnx_opt --dict ppocr_keys_v1.txt --ep auto --use_cls

import os
import glob
import argparse
import time
from ocr_runtime_module import ONNXOCR, PaddleRunner

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--imgs", type=str, required=True, help="image folder or single image")
    ap.add_argument("--onnx_dir", type=str, default=None, help="Folder containing det/rec/cls onnx")
    ap.add_argument("--det", type=str, default=None, help="Path to det onnx")
    ap.add_argument("--rec", type=str, default=None, help="Path to rec onnx")
    ap.add_argument("--cls", type=str, default=None, help="Path to cls onnx (optional)")
    ap.add_argument("--dict", type=str, required=True, help="ppocr_keys_v1.txt")
    ap.add_argument("--use_cls", action="store_true", help="run angle classifier")
    ap.add_argument("--ep", type=str, default="auto", choices=["auto","cpu","cuda","trt"], help="Execution Provider preference")
    ap.add_argument("--trt_fp16", action="store_true", help="Enable TensorRT FP16 (if TRT is used)")
    ap.add_argument("--trt_ws", type=int, default=(2<<30), help="TRT workspace size bytes")
    ap.add_argument("--rec_bs", type=int, default=6, help="REC batch size")
    ap.add_argument("--bench_paddle", action="store_true", help="Also run PaddleOCR for comparison")
    return ap.parse_args()

def list_images(p):
    if os.path.isdir(p):
        imgs = []
        for ext in ["*.jpg","*.jpeg","*.png","*.bmp","*.webp","*.tif","*.tiff"]:
            imgs.extend(glob.glob(os.path.join(p, ext)))
        return sorted(imgs)
    else:
        return [p]

def main():
    args = parse_args()

    onnx_ocr = ONNXOCR(
        onnx_dir=args.onnx_dir,
        det_path=args.det,
        rec_path=args.rec,
        cls_path=args.cls,
        dict_path=args.dict,
        ep=args.ep,
        trt_fp16=args.trt_fp16,
        trt_workspace=args.trt_ws,
        rec_bs=args.rec_bs,
        use_warp_crop=True
    )

    print(f"[EP avail] {onnx_ocr.avail}")
    print(f"[EP chosen] {onnx_ocr.providers_det}")

    imgs = list_images(args.imgs)

    # 옵션: Paddle 비교 세팅
    paddle_runner = None
    if args.bench_paddle:
        try:
            paddle_runner = PaddleRunner(use_angle_cls=args.use_cls, lang='ch')
        except Exception as e:
            print("[WARN] PaddleOCR init failed:", e)

    for path in imgs:
        print(f"\n[TEST] image: {path}")
        import cv2
        img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            print("  ! skip (cannot read)")
            continue

        # ONNX warmup (간단 1회)
        try:
            onnx_ocr.run(img, use_cls=args.use_cls, rec_bs=args.rec_bs)
        except Exception:
            pass

        t0 = time.perf_counter()
        boxes, recs, T = onnx_ocr.run(img, use_cls=args.use_cls, rec_bs=args.rec_bs)
        t1 = time.perf_counter()

        print("\n[ONNX] results:")
        for (txt, conf) in recs[:10]:
            print(f" - {txt} ({conf:.3f})")
        print(f"[ONNX] time: {(t1-t0)*1000.0:.1f} ms  | T={T}")

        if paddle_runner:
            try:
                pres, p_ms = paddle_runner.run(img)
                print("\n[Paddle] results:")
                for (txt, conf) in pres[:10]:
                    print(f" - {txt} ({conf:.3f})")
                print(f"[Paddle] time: {p_ms:.1f} ms")
            except Exception as e:
                print("[Paddle] failed:", e)

if __name__ == "__main__":
    import numpy as np  # needed by imdecode
    main()