# -*- coding: utf-8 -*-
"""
腾讯天御点选验证码 OCR 识别引擎集合 (可插拔 + 集成投票)

设计动机
--------
单一 OCR 对生僻字(如 "筹")置信度极低:  ddddocr 对 "筹" 的 CTC 最大概率仅 0.005,
而另一侧 PP-OCRv4 对 "雹" 又会误识为 "电"。两者互补, 因此本模块做 **概率级集成**:
对同一个字形, 把每个引擎在每个目标字上的 CTC 最大时间步概率加权求和, 再交给
全局指派算法, 可在 3 选 1 的题型上取得远高于单引擎的准确率。

引擎统一接口
------------
    probabilities(gray: np.ndarray) -> (chars: List[str], probs: np.ndarray)  # probs: (T, C)

gray 约定: 白底黑字的 uint8 灰度图 (字形像素值低)。
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import cv2
import numpy as np

logger = logging.getLogger("tcaptcha_ocr")

MODEL_DIR = Path(__file__).resolve().parent / "models"
PPOCR_DIR = MODEL_DIR / "ppocr"


class BaseOcrEngine:
    name = "base"

    def probabilities(self, gray: np.ndarray) -> Tuple[List[str], np.ndarray]:
        raise NotImplementedError

    def close(self):
        pass


class DdddOcrEngine(BaseOcrEngine):
    """ddddocr (CTC, 8210 字集) —— 对常用字准确率高, 集成中的稳健票"""

    name = "ddddocr"

    def __init__(self, beta: bool = False, ocr=None):
        self._ocr = ocr
        self._beta = beta

    @property
    def ocr(self):
        if self._ocr is None:
            import ddddocr

            self._ocr = ddddocr.DdddOcr(show_ad=False, beta=self._beta)
        return self._ocr

    def probabilities(self, gray: np.ndarray) -> Tuple[List[str], np.ndarray]:
        res = self.ocr.classification(gray, probability=True)
        probs = np.asarray(res["probabilities"])  # (T, 1, C)
        if probs.ndim == 3:
            probs = probs[:, 0, :]
        return list(res["charset"]), probs


class PpOcrEngine(BaseOcrEngine):
    """PaddleOCR PP-OCRv4 中文识别模型 (ONNX, 6623 字集) —— 对生僻字覆盖更好"""

    name = "ppocr_v4"

    def __init__(self, model_path: Optional[Path] = None, keys_path: Optional[Path] = None,
                 rec_height: int = 48):
        import onnxruntime as ort

        self.model_path = Path(model_path or PPOCR_DIR / "ch_PP-OCRv4_rec_infer.onnx")
        self.keys_path = Path(keys_path or PPOCR_DIR / "ppocr_keys_v1.txt")
        if not self.model_path.exists():
            raise FileNotFoundError(
                f"缺少 PP-OCRv4 识别模型: {self.model_path}\n"
                "下载: https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx"
            )
        if not self.keys_path.exists():
            raise FileNotFoundError(f"缺少 PP-OCR 字典: {self.keys_path}")

        so = ort.SessionOptions()
        so.intra_op_num_threads = 2
        so.inter_op_num_threads = 2
        self.sess = ort.InferenceSession(str(self.model_path), so, providers=["CPUExecutionProvider"])
        self.input_name = self.sess.get_inputs()[0].name
        self.rec_height = rec_height

        keys = self.keys_path.read_text(encoding="utf-8").split("\n")
        while keys and keys[-1] == "":
            keys.pop()
        # PaddleOCR 约定: 输出维度 = 1(blank) + 字典 + 1(空格)
        self.chars = ["blank"] + keys + [" "]

    def _preprocess(self, gray: np.ndarray) -> np.ndarray:
        g = gray.astype(np.float32) / 255.0
        ratio = self.rec_height / float(g.shape[0])
        g = cv2.resize(g, (0, 0), fx=ratio, fy=ratio, interpolation=cv2.INTER_LINEAR)
        g = (g - 0.5) / 0.5
        g = g[None, None, :, :]
        return np.concatenate([g, g, g], axis=1).astype(np.float32)

    def probabilities(self, gray: np.ndarray) -> Tuple[List[str], np.ndarray]:
        logits = self.sess.run(None, {self.input_name: self._preprocess(gray)})[0]  # (1, T, C)
        return self.chars, np.asarray(logits[0])


class EnsembleOcr:
    """概率级加权集成: score(t) = Σ w_i * max_t P_i(t)"""

    def __init__(self, engines: Sequence[BaseOcrEngine], weights: Optional[Sequence[float]] = None):
        self.engines = list(engines)
        self.weights = list(weights) if weights else [1.0] * len(self.engines)
        if len(self.weights) != len(self.engines):
            raise ValueError("engines 与 weights 数量不一致")

    @classmethod
    def default(cls) -> "EnsembleOcr":
        engines: List[BaseOcrEngine] = []
        try:
            engines.append(PpOcrEngine())
        except FileNotFoundError as exc:
            logger.warning("PP-OCRv4 不可用, 退化为单引擎: %s", exc)
        engines.append(DdddOcrEngine())
        return cls(engines)

    def per_engine_scores(self, gray: np.ndarray, targets: Sequence[str]) -> List[Dict[str, float]]:
        per: List[Dict[str, float]] = []
        for engine in self.engines:
            try:
                chars, probs = engine.probabilities(gray)
            except Exception as exc:  # 单个引擎异常不影响整体求解
                logger.warning("引擎 %s 推理失败: %s", engine.name, exc)
                per.append({t: 0.0 for t in targets})
                continue
            index = {c: i for i, c in enumerate(chars)}
            row = {}
            for t in targets:
                j = index.get(t)
                row[t] = float(probs[:, j].max()) if j is not None else 0.0
            per.append(row)
        return per

    def scores(self, gray: np.ndarray, targets: Sequence[str]) -> Dict[str, float]:
        per = self.per_engine_scores(gray, targets)
        out = {}
        for t in targets:
            out[t] = float(sum(w * r.get(t, 0.0) for w, r in zip(self.weights, per)))
        return out

    def close(self):
        for engine in self.engines:
            engine.close()


if __name__ == "__main__":
    import sys

    ens = EnsembleOcr.default()
    img = cv2.imread(sys.argv[1], cv2.IMREAD_GRAYSCALE)
    targets = sys.argv[2:]
    print(ens.per_engine_scores(img, targets))
    print("ensemble:", ens.scores(img, targets))