# -*- coding: utf-8 -*-
"""pokemon OCR step2: 对 images_pokemon_out/crops 跑 PaddleOCR, 输出分行文本。
复用 grading_ocr_step2 逻辑, 路径改为 images_pokemon_out。"""
import os, sys, json, csv, time
sys.stdout.reconfigure(encoding="utf-8")

OUT = "/home/martin/顾工交接/wzj/images_pokemon_out"
RECORDS = os.path.join(OUT, "crop_records.json")

print("[OCR] 初始化 PaddleOCR ...")
from paddleocr import PaddleOCR
_t_load0 = time.perf_counter()
ocr = PaddleOCR(use_doc_orientation_classify=True, use_doc_unwarping=False,
                use_textline_orientation=True)
t_load = time.perf_counter() - _t_load0
print(f"[OCR] 加载完成 (耗时 {t_load:.2f}s)\n")

with open(RECORDS, encoding="utf-8") as f:
    records = json.load(f)

ocr_results = {}
rows_csv = []
t_ocr = 0.0          # OCR 推理累计
n_ok = 0
print(f"[OCR] 处理 {len(records)} 张裁剪图 ...\n")
_t_all0 = time.perf_counter()
for i, rec in enumerate(records):
    name = rec["file"]
    crop_path = rec["crop_path"]
    try:
        _t = time.perf_counter()
        result = ocr.predict(crop_path)
        t_ocr += time.perf_counter() - _t
        n_ok += 1
    except Exception as e:
        ocr_results[name] = {"lines": [], "error": str(e)}
        rows_csv.append([name, rec["true_label"], str(rec.get("box") or ""), rec.get("conf") or "", ""])
        if (i+1) % 100 == 0:
            print(f"  [{i+1}/{len(records)}] (有错误)", flush=True)
        continue
    lines = []
    for res in result:
        if isinstance(res, dict):
            t = res.get("rec_texts")
            if t: lines.extend(list(t))
        elif hasattr(res, "rec_texts") and res.rec_texts:
            lines.extend(list(res.rec_texts))
    text_joined = " | ".join(lines) if lines else ""
    ocr_results[name] = {"lines": lines, "text": text_joined}
    rows_csv.append([name, rec["true_label"], str(rec.get("box") or ""), rec.get("conf") or "", text_joined])
    if (i+1) % 200 == 0:
        print(f"  [{i+1}/{len(records)}] 完成", flush=True)

with open(os.path.join(OUT, "ocr_result.json"), "w", encoding="utf-8") as f:
    json.dump(ocr_results, f, ensure_ascii=False, indent=2)
with open(os.path.join(OUT, "ocr_summary.csv"), "w", encoding="utf-8-sig", newline="") as f:
    w = csv.writer(f)
    w.writerow(["filename", "true_label", "box_xyxy", "conf", "ocr_text_lines"])
    w.writerows(rows_csv)

n_empty = sum(1 for v in ocr_results.values() if not v.get("lines"))
print(f"\n[OCR] 完成. JSON + CSV 已保存. 共 {len(records)} 张, OCR空文本 {n_empty} 张")

# ===== 计时报告 =====
t_all = time.perf_counter() - _t_all0
N = max(len(records), 1)
print("\n" + "=" * 50)
print("⏱  计时报告 (pokemon_ocr_step2.py)")
print("=" * 50)
print(f"  PaddleOCR 加载        : {t_load:7.2f}s")
print(f"  OCR 推理(累计)        : {t_ocr:7.2f}s  | 均 {t_ocr/N*1000:6.1f} ms/张")
print(f"  {'-'*46}")
print(f"  处理总耗时({N}张)     : {t_all:7.2f}s  | 均 {t_all/N*1000:6.1f} ms/张  | 吞吐 {N/t_all:.2f} 张/s")
print("=" * 50)

# 写 timing.log
import datetime
with open(os.path.join(OUT, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === 步骤② pokemon_ocr_step2.py (PaddleOCR 评级标签识别) ===\n")
    lf.write(f"  样本数: {N} 张 (OCR空文本 {n_empty})\n")
    lf.write(f"  PaddleOCR 加载            : {t_load:.2f}s\n")
    lf.write(f"  OCR 推理累计              : {t_ocr:.2f}s  (均 {t_ocr/N*1000:.1f} ms/张)\n")
    lf.write(f"  处理总耗时                : {t_all:.2f}s  (均 {t_all/N*1000:.1f} ms/张, 吞吐 {N/t_all:.2f} 张/s)\n")
print(f"✓ 计时已写入: {os.path.join(OUT, 'timing.log')}")