"""Reference of the browser OCR pipeline (src/lib/docai/ocr.js) on PP-OCRv5 mobile det + en_PP-OCRv5 mobile rec.

  det: longest side resized to LIMIT (multiple of 32), BGR, ImageNet mean/std -> probability map;
       threshold 0.3 -> connected components -> axis-aligned boxes scored by mean probability (>= 0.6),
       grown by area * 1.5 / perimeter (DB's unclip for a rectangle)
  rec: each box cropped, height 48, width bucketed, (x/255 - 0.5) / 0.5, BGR -> CTC greedy decode;
       each kept character's time step gives its x position, so lines split into word boxes.
Run with .models/.venv311mp (opencv, onnxruntime). usage: python ocr_reference.py [image]
"""
import os, sys, time
import numpy as np
import cv2
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont

HERE = os.path.dirname(os.path.abspath(__file__))
PORT = os.path.join(HERE, "..", "portfolio", "pp-ocrv5")
DICT = open(os.path.join(PORT, "en_dict.txt"), encoding="utf-8").read().split("\n")[:-1] + [" "]
LIMIT, THRESH, BOX_THRESH, UNCLIP = 1280, 0.3, 0.6, 1.5
MEAN, STD = np.array([0.485, 0.456, 0.406], np.float32), np.array([0.229, 0.224, 0.225], np.float32)


def synthetic_invoice():
    img = Image.new("RGB", (1240, 1600), "white")
    d = ImageDraw.Draw(img)
    f = lambda s, b=False: ImageFont.truetype("C:/Windows/Fonts/arialbd.ttf" if b else "C:/Windows/Fonts/arial.ttf", s)
    d.text((80, 80), "ACME Supplies Ltd.", font=f(44, True), fill="black")
    d.text((900, 90), "INVOICE", font=f(52, True), fill=(240, 72, 44))
    d.text((80, 150), "221B Baker Street, London NW1 6XE", font=f(24), fill="#333")
    for i, (k, v) in enumerate([("Invoice #:", "INV-2024-0412"), ("Date:", "12 March 2024"), ("Due:", "11 April 2024")]):
        d.text((820, 200 + i * 36), k, font=f(24, True), fill="black")
        d.text((960, 200 + i * 36), v, font=f(24), fill="black")
    d.text((80, 360), "Description", font=f(26, True), fill="black")
    d.text((900, 360), "Amount", font=f(26, True), fill="black")
    for i, (k, v) in enumerate([("Printer paper A4 (10 reams)", "$45.90"), ("Toner cartridge HP 26A", "$129.00"), ("Delivery", "$12.50")]):
        d.text((80, 420 + i * 48), k, font=f(26), fill="black")
        d.text((900, 420 + i * 48), v, font=f(26), fill="black")
    d.text((700, 620), "Total due: $187.40", font=f(34, True), fill="black")
    d.text((80, 760), "Thank you for your business! Questions? billing@acme.co.uk", font=f(20), fill="#555")
    return np.array(img)[:, :, ::-1]  # BGR


def detect(sess, bgr):
    h, w = bgr.shape[:2]
    s = min(1.0, LIMIT / max(h, w)) if max(h, w) > LIMIT else LIMIT / max(h, w)
    rh, rw = max(32, int(round(h * s / 32)) * 32), max(32, int(round(w * s / 32)) * 32)
    x = cv2.resize(bgr, (rw, rh)).astype(np.float32) / 255
    x = ((x - MEAN) / STD).transpose(2, 0, 1)[None]
    prob = sess.run(None, {"x": x})[0][0, 0]
    n, labels, stats, _ = cv2.connectedComponentsWithStats((prob > THRESH).astype(np.uint8), connectivity=8)
    boxes = []
    for i in range(1, n):
        x0, y0, bw, bh, area = stats[i]
        if min(bw, bh) < 3:
            continue
        score = float(prob[labels == i].mean())
        if score < BOX_THRESH:
            continue
        d = bw * bh * UNCLIP / (2 * (bw + bh))
        sx, sy = w / rw, h / rh
        boxes.append([max(0, (x0 - d) * sx), max(0, (y0 - d) * sy), min(w, (x0 + bw + d) * sx), min(h, (y0 + bh + d) * sy), score])
    boxes.sort(key=lambda b: (round(b[1] / 20), b[0]))
    return boxes


def recognize(sess, bgr, boxes):
    lines = []
    for x0, y0, x1, y1, _ in boxes:
        crop = bgr[int(y0):int(np.ceil(y1)), int(x0):int(np.ceil(x1))]
        ch, cw = crop.shape[:2]
        tw = max(16, int(np.ceil(48 * cw / ch)))
        x = cv2.resize(crop, (tw, 48)).astype(np.float32) / 255
        x = ((x - 0.5) / 0.5).transpose(2, 0, 1)[None]
        probs = sess.run(None, {"x": x})[0][0]  # [T, 438]
        T = probs.shape[0]
        ids, conf = probs.argmax(1), probs.max(1)
        chars, prev = [], 0
        for t, k in enumerate(ids):
            if k != 0 and k != prev:
                chars.append((DICT[k - 1], t, conf[t]))
            prev = k
        text = "".join(c for c, _, _ in chars)
        # words from the character time steps
        words, cur = [], []
        for c in chars + [(" ", T, 1.0)]:
            if c[0] == " ":
                if cur:
                    t0, t1 = cur[0][1], cur[-1][1]
                    words.append(("".join(k for k, _, _ in cur), x0 + (t0 - 0.5) * cw / T, x0 + (t1 + 1.5) * cw / T, float(np.mean([p for _, _, p in cur]))))
                cur = []
            else:
                cur.append(c)
        lines.append((text, float(np.mean([p for _, _, p in chars])) if chars else 0, (x0, y0, x1, y1), words))
    return lines


if __name__ == "__main__":
    bgr = cv2.imread(sys.argv[1]) if len(sys.argv) > 1 else synthetic_invoice()
    det = ort.InferenceSession(os.path.join(PORT, "det.onnx"), providers=["CPUExecutionProvider"])
    rec = ort.InferenceSession(os.path.join(PORT, "rec.onnx"), providers=["CPUExecutionProvider"])
    t0 = time.time()
    boxes = detect(det, bgr)
    lines = recognize(rec, bgr, boxes)
    print(f"{len(lines)} lines in {time.time() - t0:.2f}s")
    for text, conf, box, words in lines:
        print(f"{conf:.2f} {[int(v) for v in box]} {text!r} | words: {[(w[0], int(w[1]), int(w[2])) for w in words]}")