"""S5：合成扫描件上的字错误率（CER）。文字页渲染成图片 → 加扫描退化 → OCR → 和文字层比。

用法：.venv/bin/python -I synth_cer.py <pdf> <页码范围> [dpi]
退化三档：clean（只渲染）、scan（灰度、倾斜 1°、轻微模糊、噪点、JPEG 60）、bad（150 dpi 后放大、倾斜 2°、更模糊、JPEG 35）。
文字层按页的字符顺序，比较时去掉空白；OCR 行按坐标从上到下、从左到右拼接。顺序不同会抬高 CER，所以同时报「字符包」差异（不看顺序）。
"""

from __future__ import annotations

import io
import random
import sys
import time
from collections import Counter

import numpy as np
import pdfplumber
import pypdfium2 as pdfium
from PIL import Image, ImageFilter
from rapidfuzz.distance import Levenshtein
from rapidocr import RapidOCR



def parse_ranges(s: str) -> set[int]:
    out: set[int] = set()
    for part in s.split(","):
        a, _, b = part.partition("-")
        out.update(range(int(a), int(b or a) + 1))
    return out


def degrade(img: Image.Image, level: str, rnd: random.Random) -> Image.Image:
    if level == "clean":
        return img
    g = img.convert("L")
    angle = rnd.uniform(-1, 1) * (1 if level == "scan" else 2)
    if level == "bad":
        g = g.resize((g.width * 3 // 4, g.height * 3 // 4)).resize((g.width, g.height))
    g = g.rotate(angle, expand=False, fillcolor=255, resample=Image.BICUBIC)
    g = g.filter(ImageFilter.GaussianBlur(0.6 if level == "scan" else 1.1))
    a = np.asarray(g).astype(np.int16)
    noise = np.random.default_rng(rnd.randint(0, 1 << 30)).integers(-25, 26, a.shape)
    a = np.clip(a + noise, 0, 255).astype(np.uint8)
    buf = io.BytesIO()
    Image.fromarray(a).save(buf, "JPEG", quality=60 if level == "scan" else 35)
    return Image.open(io.BytesIO(buf.getvalue())).convert("RGB")


def norm(s: str) -> str:
    return "".join(s.split()).replace("（", "(").replace("）", ")").replace("：", ":").replace("，", ",").replace("；", ";")


def main() -> None:
    pdf_path, pages = sys.argv[1], sorted(parse_ranges(sys.argv[2]))
    dpi = int(sys.argv[3]) if len(sys.argv) > 3 else 200
    eng = RapidOCR()
    doc = pdfium.PdfDocument(pdf_path)
    rnd = random.Random(7)
    stats = {lv: [0, 0, 0, 0.0] for lv in ("clean", "scan", "bad")}  # 编辑距离、字符包差、总字数、秒
    with pdfplumber.open(pdf_path) as pl:
        for no in pages:
            page = pl.pages[no - 1]
            truth = norm("".join(c["text"] for c in page.chars))
            if len(truth) < 100:
                continue
            base = doc[no - 1].render(scale=dpi / 72).to_pil()
            for lv in stats:
                img = degrade(base, lv, rnd)
                t0 = time.time()
                r = eng(np.asarray(img))
                dt = time.time() - t0
                items = sorted(zip(r.boxes, r.txts), key=lambda x: (round(x[0][0][1] / 20), x[0][0][0])) if r.boxes is not None else []
                got = norm("".join(t for _, t in items))
                bag = sum(((Counter(truth) - Counter(got)) + (Counter(got) - Counter(truth))).values())
                st = stats[lv]
                st[0] += Levenshtein.distance(truth, got)
                st[1] += bag
                st[2] += len(truth)
                st[3] += dt
            print(f"p{no} {len(truth)} chars", flush=True)
    n = sum(1 for _ in pages)
    for lv, (ed, bag, total, sec) in stats.items():
        print(f"{lv:6s} CER(顺序) {ed / total:.2%}  字符包差 {bag / total / 2:.2%}  共 {total} 字  OCR {sec:.1f}s")
    print(f"pages tried {n}, dpi {dpi}")


if __name__ == "__main__":
    main()