"""paddleocr.dev CPU benchmark: one model config per process.

Usage: python bench.py --engine onnxruntime --tier PP-OCRv6_medium [--mkldnn 0|1] [--runs 3]
Prints one JSON object on the last line.
"""
import argparse, json, os, platform, resource, statistics, sys, time
from pathlib import Path

HERE = Path(__file__).parent
SAMPLES = HERE / "samples"

RECEIPT = ["HARBOR STREET COFFEE", "214 Harbor Street, Portland ME", "Tel (207) 555-0148",
           "Order #4471", "2026-10-21 08:42", "Cashier: Dana", "Table 6",
           "2 x Flat White", "9.50", "1 x Oat Latte", "5.25", "1 x Blueberry Muffin", "3.75",
           "1 x Avocado Toast", "11.00", "1 x Sparkling Water", "2.50",
           "Subtotal", "32.00", "Sales tax 5.5%", "1.76", "TOTAL", "$33.76", "VISA **** 4012", "33.76",
           "Thank you! Come again.", "Free Wi-Fi: HarborGuest"]
GT = {
    "receipt.png": RECEIPT,
    "receipt-photo.jpg": RECEIPT,
    "invoice.png": ["Northwind Lighting Co.", "88 Foundry Lane, Leeds LS1 4AB, United Kingdom",
                    "billing@northwind-lighting.example · VAT GB 123 4567 89", "INVOICE",
                    "Invoice no.", "INV-2026-0917", "Issue date", "14 Oct 2026", "Due date", "13 Nov 2026",
                    "BILL TO", "Ashgrove Architects Ltd", "Unit 5, Canal Wharf", "Manchester M1 2EQ",
                    "PROJECT", "Library refit, level 2", "PO number: AA-77310",
                    "Description", "Qty", "Unit price", "Amount",
                    "LED panel 600x600, 4000K", "24", "£46.00", "£1,104.00",
                    "Linear pendant, 1.2 m", "10", "£128.50", "£1,285.00",
                    "Emergency exit sign", "6", "£39.90", "£239.40",
                    "DALI dimming driver", "12", "£22.75", "£273.00",
                    "Installation labour (hours)", "16", "£55.00", "£880.00",
                    "Subtotal", "£3,781.40", "VAT 20%", "£756.28", "Total due", "£4,537.68",
                    "Payment terms",
                    "Please pay within 30 days by bank transfer to Northwind Lighting Co.",
                    "Sort code 12-34-56, account 87654321. Quote the invoice number as reference.",
                    "Late payments may incur interest at 8% above the Bank of England base rate.",
                    "Thank you for your business."],
    "notice-zh-en.png": ["城东社区图书馆", "Chengdong Community Library",
                         "开放时间：周一至周五 9:00–21:00", "Opening hours: Mon–Fri 9:00–21:00",
                         "周末及法定节假日 10:00–17:00", "Weekends and public holidays 10:00–17:00",
                         "请保持安静，馆内禁止饮食。", "Please keep quiet. No food or drink inside.",
                         "借书证办理请到一楼服务台。", "Library cards are issued at the ground-floor desk."],
    "notice-ja.png": ["臨時休業のお知らせ", "いつもご利用いただき誠にありがとうございます。",
                      "設備点検のため、下記の日程で臨時休業いたします。", "休業日：2026年11月18日（水）",
                      "営業再開：2026年11月19日（木）午前10時", "お客様にはご迷惑をおかけしますが、",
                      "何卒ご理解のほどよろしくお願い申し上げます。", "みなと町ベーカリー　店主"],
}

# Character normalisation applied to BOTH sides before scoring, so we score
# recognition rather than typographic choices (full-width vs ASCII punctuation, dashes, spaces).
TRANS = str.maketrans({"：": ":", "，": ",", "。": ".", "（": "(", "）": ")", "–": "-", "—": "-", "　": " ", "·": "·"})


def norm(s):
    return "".join(s.translate(TRANS).split())


def lev(a, b):
    if a == b:
        return 0
    if len(a) < len(b):
        a, b = b, a
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        cur = [i]
        for j, cb in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
        prev = cur
    return prev[-1]


def score(gt_lines, pred_lines):
    """Character error rate over ground-truth items.

    Each GT item is matched to the predicted line that contains it best
    (substring match counts as exact, since detectors may merge neighbouring
    items onto one line). Unmatched GT items cost their full length."""
    preds = [norm(p) for p in pred_lines]
    total = edits = exact = 0
    for g in gt_lines:
        ng = norm(g)
        total += len(ng)
        best = len(ng)
        for p in preds:
            if ng in p:
                best = 0
                break
            best = min(best, lev(ng, p) if abs(len(p) - len(ng)) < len(ng) * 2 + 4 else best)
        edits += best
        exact += best == 0
    return {"cer": round(edits / total, 4), "exact_items": exact, "items": len(gt_lines)}


def dir_size(p):
    return sum(f.stat().st_size for f in Path(p).rglob("*") if f.is_file())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--engine", default="paddle")
    ap.add_argument("--tier", default="PP-OCRv6_medium")
    ap.add_argument("--mkldnn", type=int, default=-1)
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--images", default="")
    ap.add_argument("--dump", default="")
    a = ap.parse_args()

    from paddleocr import PaddleOCR
    import paddleocr
    kw = dict(use_doc_orientation_classify=False, use_doc_unwarping=False, use_textline_orientation=False,
              text_detection_model_name=f"{a.tier}_det", text_recognition_model_name=f"{a.tier}_rec")
    if a.engine != "paddle":
        kw["engine"] = a.engine
    if a.mkldnn >= 0:
        kw["enable_mkldnn"] = bool(a.mkldnn)
    t0 = time.perf_counter()
    ocr = PaddleOCR(**kw)
    init_s = time.perf_counter() - t0

    images = [i for i in GT] if not a.images else a.images.split(",")
    per = {}
    for name in images:
        path = str(SAMPLES / name)
        t = time.perf_counter()
        res = ocr.predict(path)[0]           # warm-up run (not timed in median)
        first = time.perf_counter() - t
        times = []
        for _ in range(a.runs):
            t = time.perf_counter()
            res = ocr.predict(path)[0]
            times.append(time.perf_counter() - t)
        texts = list(res["rec_texts"])
        scores = [float(s) for s in res["rec_scores"]]
        per[name] = {"first_s": round(first, 3), "median_s": round(statistics.median(times), 3),
                     "runs_s": [round(x, 3) for x in times], "lines": len(texts),
                     "mean_conf": round(sum(scores) / max(1, len(scores)), 4),
                     **score(GT[name], texts)}
        if a.dump:
            out = Path(a.dump) / f"{a.engine}-{a.tier}-{name}.json"
            out.parent.mkdir(parents=True, exist_ok=True)
            data = {"rec_texts": texts, "rec_scores": scores,
                    "rec_polys": [p.tolist() for p in res["rec_polys"]],
                    "rec_boxes": [b.tolist() for b in res["rec_boxes"]]}
            out.write_text(json.dumps(data, ensure_ascii=False))

    models_dir = Path.home() / ".paddlex" / "official_models"
    suffix = "_onnx" if a.engine == "onnxruntime" else ""
    sizes = {}
    for part in ("det", "rec"):
        d = models_dir / f"{a.tier}_{part}{suffix}"
        sizes[part] = dir_size(d) if d.exists() else None
    try:
        import paddle
        pv = paddle.__version__
    except Exception:
        pv = None
    out = {"engine": a.engine, "tier": a.tier, "mkldnn": a.mkldnn, "init_s": round(init_s, 2),
           "peak_rss_mb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 1),
           "model_bytes": sizes, "paddleocr": paddleocr.__version__, "paddle": pv,
           "python": platform.python_version(), "cpus": os.cpu_count(), "per_image": per}
    print(json.dumps(out, ensure_ascii=False))


if __name__ == "__main__":
    main()