#!/usr/bin/env python3
"""
PP-OCRv6 ONNX Inference (standalone, zero Paddle dependency)

Dependencies:
    numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper

Usage:
    from ppocrv6_onnx import PPOCRv6Onnx
    ocr = PPOCRv6Onnx(det_onnx="det.onnx", rec_onnx="rec.onnx", char_dict="inference.yml")
    results = ocr(img)  # img is BGR numpy array

    # With direction classifier:
    ocr = PPOCRv6Onnx(..., cls_onnx="cls.onnx", cls_label_list=["0","180"], cls_thresh=0.9)
"""

import argparse
import json
import math
import os
from typing import List, Optional, Tuple, Union

import cv2
import numpy as np
import axengine as ort
import yaml

from PIL import Image, ImageDraw, ImageFont
from shapely.geometry import Polygon
import pyclipper
import random


# ============================================================================
# Helpers
# ============================================================================

def _get_dim_value(dim):
    """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic."""
    if dim is None:
        return 0
    if isinstance(dim, str):
        return 0
    if hasattr(dim, "dim_value"):
        return int(dim.dim_value) if dim.dim_value else 0
    if hasattr(dim, "dim_param"):
        return 0
    try:
        return int(dim)
    except (TypeError, ValueError):
        return 0


def _detect_fixed_dims(session: ort.InferenceSession, det_onnx: str):
    """Detect fixed H/W from ONNX input shape. Returns (fixed_h, fixed_w)."""
    inp = session.get_inputs()[0]
    h, w = _get_dim_value(inp.shape[2]), _get_dim_value(inp.shape[3])
    if h == 0 and w == 0:
        try:
            import onnx
            m = onnx.load(det_onnx)
            dims = m.graph.input[0].type.tensor_type.shape.dim
            h = dims[2].dim_value if len(dims) > 2 else 0
            w = dims[3].dim_value if len(dims) > 3 else 0
        except Exception:
            pass
    return (h if h > 0 else 0), (w if w > 0 else 0)


def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
    if isinstance(source, list):
        return source
    ext = os.path.splitext(source)[1].lower()
    if ext in (".yml", ".yaml"):
        with open(source, "r", encoding="utf-8") as f:
            cfg = yaml.safe_load(f)
        dic = cfg.get("PostProcess", {}).get("character_dict", [])
        if not dic:
            raise ValueError(f"No PostProcess.character_dict found in {source}")
        return dic
    elif ext == ".txt":
        with open(source, "r", encoding="utf-8") as f:
            return [line.strip("\n\r") for line in f.readlines()]
    else:
        raise ValueError(f"Unsupported char_dict source: {source}. Use .yml, .txt, or list.")


# ============================================================================
# 1. Detection Preprocessing
# ============================================================================

class _DetResizeForTest:
    def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000):
        self.limit_side_len = limit_side_len
        self.limit_type = limit_type
        self.max_side_limit = max_side_limit

    def _image_padding(self, im, value=0):
        h, w, c = im.shape
        im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
        im_pad[:h, :w, :] = im
        return im_pad

    def _resize_image_type0(self, img):
        h, w, _ = img.shape
        limit_side_len = self.limit_side_len
        if self.limit_type == "max":
            ratio = float(limit_side_len) / max(h, w) if max(h, w) > limit_side_len else 1.0
        elif self.limit_type == "min":
            ratio = float(limit_side_len) / min(h, w) if min(h, w) < limit_side_len else 1.0
        elif self.limit_type == "resize_long":
            ratio = float(limit_side_len) / max(h, w)
        else:
            raise ValueError(f"not support limit_type: {self.limit_type}")

        resize_h, resize_w = int(h * ratio), int(w * ratio)
        if max(resize_h, resize_w) > self.max_side_limit:
            ratio = float(self.max_side_limit) / max(resize_h, resize_w)
            resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)
        resize_h = max(int(round(resize_h / 32) * 32), 32)
        resize_w = max(int(round(resize_w / 32) * 32), 32)
        if int(resize_w) <= 0 or int(resize_h) <= 0:
            return None, (None, None)
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
        ratio_h, ratio_w = resize_h / float(h), resize_w / float(w)
        return img, [ratio_h, ratio_w]

    def __call__(self, img):
        src_h, src_w = img.shape[:2]
        if sum([src_h, src_w]) < 64:
            img = self._image_padding(img)
        img, [ratio_h, ratio_w] = self._resize_image_type0(img)
        shape = np.array([src_h, src_w, ratio_h, ratio_w])
        return img, shape


class _NormalizeImage:
    def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
        self.scale = np.float32(scale)
        shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
        self.mean = np.array(mean, dtype=np.float32).reshape(shape)
        self.std = np.array(std, dtype=np.float32).reshape(shape)

    def __call__(self, img):
        return (img.astype("float32") * self.scale - self.mean) / self.std


class _ToCHWImage:
    def __call__(self, img):
        return img.transpose((2, 0, 1))


# ============================================================================
# 2. Recognition Preprocessing
# ============================================================================

def _resize_norm_img(img, image_shape, max_wh_ratio=None):
    imgC, imgH, imgW = image_shape
    if max_wh_ratio is None:
        max_wh_ratio = imgW * 1.0 / imgH
        h, w = img.shape[:2]
        max_wh_ratio = max(max_wh_ratio, w / h)
    target_w = int(imgH * max_wh_ratio)
    h, w = img.shape[:2]
    ratio = w / h
    resized_w = target_w if math.ceil(imgH * ratio) > target_w else int(math.ceil(imgH * ratio))
    resized_image = cv2.resize(img, (resized_w, imgH)).astype("float32")
    resized_image = resized_image.transpose((2, 0, 1))
    # resized_image /= 255.0
    # resized_image -= 0.5
    # resized_image /= 0.5
    padding_im = np.zeros((imgC, imgH, target_w), dtype=np.float32)
    padding_im[:, :, 0:resized_w] = resized_image
    return padding_im


# ============================================================================
# 3. Detection Postprocessing
# ============================================================================

class _DBPostProcess:
    def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000,
                 unclip_ratio=2.0, use_dilation=False, score_mode="fast", box_type="quad"):
        self.thresh = thresh
        self.box_thresh = box_thresh
        self.max_candidates = max_candidates
        self.unclip_ratio = unclip_ratio
        self.min_size = 3
        self.score_mode = score_mode
        self.box_type = box_type
        self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])

    def _unclip(self, box, unclip_ratio):
        poly = Polygon(box)
        distance = poly.area * unclip_ratio / poly.length
        offset = pyclipper.PyclipperOffset()
        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
        return offset.Execute(distance)

    def _get_mini_boxes(self, contour):
        bb = cv2.minAreaRect(contour)
        points = sorted(list(cv2.boxPoints(bb)), key=lambda x: x[0])
        i1, i4 = (0, 1) if points[1][1] > points[0][1] else (1, 0)
        i2, i3 = (2, 3) if points[3][1] > points[2][1] else (3, 2)
        return [points[i1], points[i2], points[i3], points[i4]], min(bb[1])

    def _box_score_fast(self, bitmap, _box):
        h, w = bitmap.shape[:2]
        box = _box.copy()
        xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
        xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
        ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
        ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        box[:, 0] -= xmin
        box[:, 1] -= ymin
        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]

    def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
        height, width = _bitmap.shape
        outs = cv2.findContours(
            (_bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
        )
        contours = outs[0] if len(outs) == 2 else outs[1]
        num = min(len(contours), self.max_candidates)
        boxes, scores = [], []
        for i in range(num):
            points, sside = self._get_mini_boxes(contours[i])
            if sside < self.min_size:
                continue
            points = np.array(points)
            score = self._box_score_fast(pred, points.reshape(-1, 2))
            if self.box_thresh > score:
                continue
            box = self._unclip(points, self.unclip_ratio)
            if len(box) > 1:
                continue
            box = np.array(box).reshape(-1, 1, 2)
            box, sside = self._get_mini_boxes(box)
            if sside < self.min_size + 2:
                continue
            box = np.array(box)
            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
            boxes.append(box.astype("int32"))
            scores.append(score)
        return np.array(boxes, dtype="int32"), scores

    def __call__(self, pred, shape_list):
        pred = pred[:, 0, :, :]
        segmentation = pred > self.thresh
        boxes_batch = []
        for bi in range(pred.shape[0]):
            src_h, src_w, ratio_h, ratio_w = shape_list[bi]
            mask = cv2.dilate(np.array(segmentation[bi]).astype(np.uint8),
                              self.dilation_kernel) if self.dilation_kernel is not None else segmentation[bi]
            boxes, _ = self._boxes_from_bitmap(pred[bi], mask, src_w, src_h)
            boxes_batch.append(boxes)
        return boxes_batch


# ============================================================================
# 4. Recognition Postprocessing
# ============================================================================

class _CTCLabelDecode:
    def __init__(self, character_list: List[str], use_space_char=True):
        self.character_str = list(character_list)
        if use_space_char:
            self.character_str.append(" ")
        dict_character = ["blank"] + self.character_str
        self.character = dict_character

    def decode(self, text_index, text_prob=None, is_remove_duplicate=True):
        result_list = []
        for bi in range(len(text_index)):
            sel = np.ones(len(text_index[bi]), dtype=bool)
            if is_remove_duplicate:
                sel[1:] = text_index[bi][1:] != text_index[bi][:-1]
            sel &= text_index[bi] != 0
            chars = [self.character[int(t)] for t in text_index[bi][sel]]
            conf = text_prob[bi][sel] if text_prob is not None else [1] * len(sel)
            if len(conf) == 0:
                conf = [0]
            result_list.append(("".join(chars), float(np.mean(conf))))
        return result_list

    def __call__(self, preds):
        return self.decode(preds.argmax(axis=2), preds.max(axis=2), is_remove_duplicate=True)


# ============================================================================
# 5. Image Utilities
# ============================================================================

def _get_rotate_crop_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
    assert len(points) == 4
    cw = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
    ch = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
    pts_std = np.float32([[0, 0], [cw, 0], [cw, ch], [0, ch]])
    M = cv2.getPerspectiveTransform(points.astype(np.float32), pts_std)
    dst = cv2.warpPerspective(img, M, (cw, ch), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC)
    if dst.shape[0] * 1.0 / dst.shape[1] >= 1.5:
        dst = np.rot90(dst)
    return dst


def _sorted_boxes(dt_boxes):
    if len(dt_boxes) == 0:
        return dt_boxes
    boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
    lst = list(boxes)
    for i in range(len(lst) - 1):
        for j in range(i, -1, -1):
            if abs(lst[j + 1][0][1] - lst[j][0][1]) < 10 and lst[j + 1][0][0] < lst[j][0][0]:
                lst[j], lst[j + 1] = lst[j + 1], lst[j]
            else:
                break
    return lst


def draw_ocr_result(
    img: np.ndarray,
    results: List[dict],
    font_path: str = "./fonts/simfang.ttf",
) -> np.ndarray:
    """Draw detection boxes (semi-transparent) on original image, with text list on the right side."""
    h, w = img.shape[:2]

    # --- left: original image with semi-transparent colored boxes ---
    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    overlay = Image.new("RGBA", pil_img.size, (0, 0, 0, 0))
    draw_overlay = ImageDraw.Draw(overlay)

    random.seed(0)
    for res in results:
        box = res["box"]
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 90)
        draw_overlay.polygon([tuple(p) for p in box], fill=color)

    left_img = Image.alpha_composite(pil_img.convert("RGBA"), overlay).convert("RGB")

    # --- right: white canvas with text list ---
    right_w = int(w * 0.9)
    right = Image.new("RGB", (right_w, h), (255, 255, 255))
    draw_right = ImageDraw.Draw(right)

    try:
        font = ImageFont.truetype(font_path, 14)
    except (OSError, IOError):
        font = ImageFont.load_default()

    y = 5
    gap = 18
    for i, res in enumerate(results):
        text = f"{i+1}. {res['text']}  ({res['confidence']:.3f})"
        # Color block matching the box
        random.seed(i)
        blk_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        draw_right.rectangle([5, y + 3, 15, y + 14], fill=blk_color, outline=(0, 0, 0))
        draw_right.text((20, y), text, fill=(0, 0, 0), font=font)
        y += gap

    # --- concat left + right ---
    result_img = Image.new("RGB", (w + right_w, h))
    result_img.paste(left_img, (0, 0))
    result_img.paste(right, (w, 0))

    return cv2.cvtColor(np.array(result_img), cv2.COLOR_RGB2BGR)


# ============================================================================
# 6. Main Inference Engine
# ============================================================================

class PPOCRv6Onnx:

    def __init__(
        self,
        det_onnx: str,
        rec_onnx: str,
        char_dict: Union[str, List[str]],
        # Detection params
        det_limit_side_len: int = 960,
        det_db_thresh: float = 0.2,
        det_db_box_thresh: float = 0.4,
        det_db_unclip_ratio: float = 1.4,
        det_max_candidates: int = 3000,
        # Recognition params
        rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
        rec_batch_num: int = 1,
        # Classifier params
        use_angle_cls: bool = False,
        cls_onnx: Optional[str] = None,
        cls_image_shape: Tuple[int, int, int] = (3, 48, 192),
        cls_batch_num: int = 1,
        cls_thresh: float = 0.9,
        cls_label_list: Optional[List[str]] = None,
        # Common
        drop_score: float = 0.5,
        use_gpu: bool = False,
        onnx_providers: Optional[List[str]] = None,
        resize_mode: str = "letterbox",
    ):
        assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"
        if cls_label_list is None:
            cls_label_list = ["0", "180"]

        self.rec_image_shape = rec_image_shape
        self.rec_batch_num = rec_batch_num
        self.drop_score = drop_score
        self.use_angle_cls = use_angle_cls
        self.cls_thresh = cls_thresh
        self.cls_label_list = cls_label_list
        self.cls_batch_num = cls_batch_num
        self.cls_image_shape = cls_image_shape
        self._resize_mode = resize_mode

        self.det_session = ort.InferenceSession(det_onnx)
        self.det_input_name = self.det_session.get_inputs()[0].name

        self.rec_session = ort.InferenceSession(rec_onnx)
        self.rec_input_name = self.rec_session.get_inputs()[0].name

        # Classifier session
        if use_angle_cls:
            if cls_onnx is None:
                raise ValueError("cls_onnx is required when use_angle_cls=True")
            self.cls_session = ort.InferenceSession(cls_onnx)
            self.cls_input_name = self.cls_session.get_inputs()[0].name
            # Detect fixed dims for cls
            cls_h, cls_w = _detect_fixed_dims(self.cls_session, cls_onnx)
            self._cls_fixed_h = cls_h if cls_h > 0 else 0
            self._cls_fixed_w = cls_w if cls_w > 0 else 0
        else:
            self.cls_session = None

        # Detect fixed dims for det/rec
        self._det_fixed_h, self._det_fixed_w = _detect_fixed_dims(self.det_session, det_onnx)
        det_shape = self.det_session.get_inputs()[0].shape
        print(f"[PPOCRv6] det shape={det_shape}, fixed_h={self._det_fixed_h}, fixed_w={self._det_fixed_w}, "
              f"cls={use_angle_cls}, resize_mode={resize_mode}")

        # Fixed rec width from ONNX
        rec_inp = self.rec_session.get_inputs()[0]
        rec_fw = _get_dim_value(rec_inp.shape[3])
        self._rec_fixed_w = rec_fw if rec_fw > 0 else 0

        # ---- Detection pre/post ----
        self._det_resize = _DetResizeForTest(limit_side_len=det_limit_side_len, limit_type="max")
        # self._det_normalize = _NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self._det_normalize = _NormalizeImage(mean=[0., 0., 0.], std=[1.0, 1.0, 1.0], scale=1.0)
        self._det_to_chw = _ToCHWImage()
        self._det_post = _DBPostProcess(thresh=det_db_thresh, box_thresh=det_db_box_thresh,
                                         unclip_ratio=det_db_unclip_ratio, max_candidates=det_max_candidates, box_type="quad")

        # ---- Recognition post ----
        self._rec_post = _CTCLabelDecode(_load_char_dict(char_dict), use_space_char=True)

    # ---- Detection with static shape support ----

    def _preprocess_det(self, img: np.ndarray):
        src_h, src_w = img.shape[:2]
        fh, fw = self._det_fixed_h, self._det_fixed_w
        # print(f'fh {fh} fw {fw}')
        if fh > 0 and fw > 0:
            if self._resize_mode == "stretch":
                # Direct resize (official behavior)
                img_r = cv2.resize(img, (fw, fh))
                shape = np.array([src_h, src_w, float(fh) / src_h, float(fw) / src_w])
            else:
                # Letterbox: ratio-preserving + pad
                ratio = min(fh / src_h, fw / src_w)
                new_h = min(max(int(round(src_h * ratio / 32) * 32), 32), fh)
                new_w = min(max(int(round(src_w * ratio / 32) * 32), 32), fw)
                img_r = cv2.resize(img, (new_w, new_h))
                pad_h, pad_w = max(0, fh - new_h), max(0, fw - new_w)
                if pad_h or pad_w:
                    img_r = cv2.copyMakeBorder(img_r, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=(0, 0, 0))
                shape = np.array([src_h * fh / new_h, src_w * fw / new_w,
                                  float(new_h) / src_h, float(new_w) / src_w])
        else:
            img_r, shape = self._det_resize(img)

        img_n = self._det_normalize(img_r)
        img_c = self._det_to_chw(img_n)
        return np.expand_dims(img_c.astype(np.float32), axis=0), shape

    def _postprocess_det(self, output, shape):
        return self._det_post(output, np.expand_dims(shape, axis=0))[0]

    def detect(self, img):
        tensor, shape = self._preprocess_det(img)
        out = self.det_session.run(None, {self.det_input_name: tensor})
        return self._postprocess_det(out[0], shape)

    # ---- Classifier ----

    def _preprocess_cls(self, img_list):
        num = len(img_list)
        width_list = [im.shape[1] / float(im.shape[0]) for im in img_list]
        indices = np.argsort(np.array(width_list))
        batches, idx_maps = [], []
        for beg in range(0, num, self.cls_batch_num):
            end = min(num, beg + self.cls_batch_num)
            imgC, imgH, imgW = self.cls_image_shape

            # Override with ONNX fixed dimensions
            if self._cls_fixed_h > 0:
                imgH = self._cls_fixed_h
            if self._cls_fixed_w > 0:
                imgW = self._cls_fixed_w

            max_wh_ratio = imgW / imgH
            for ino in range(beg, end):
                h, w = img_list[indices[ino]].shape[:2]
                max_wh_ratio = max(max_wh_ratio, w / h)
            if self._cls_fixed_w > 0:
                max_wh_ratio = self._cls_fixed_w / imgH

            shape = (imgC, imgH, imgW)
            norm_list, idx_list = [], []
            for ino in range(beg, end):
                orig_idx = indices[ino]
                norm = _resize_norm_img(img_list[orig_idx], shape, max_wh_ratio=max_wh_ratio)
                norm_list.append(np.expand_dims(norm, axis=0))
                idx_list.append(orig_idx)
            if norm_list:
                batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
            idx_maps.append(idx_list)
        return batches, idx_maps

    def _postprocess_cls(self, batch_outputs, idx_maps, total_num, img_list):
        results = [("0", 1.0)] * total_num
        for preds_batch, idx_list in zip(batch_outputs, idx_maps):
            pred_ids = preds_batch.argmax(axis=1)
            for i, orig_idx in enumerate(idx_list):
                label = self.cls_label_list[int(pred_ids[i])]
                score = float(preds_batch[i, int(pred_ids[i])])
                results[orig_idx] = (label, score)
                if "180" in str(label) and score > self.cls_thresh:
                    img_list[orig_idx] = cv2.rotate(img_list[orig_idx], cv2.ROTATE_180)
        return results

    def classify(self, img_list):
        if not img_list or not self.use_angle_cls:
            return img_list, [], 0
        img_list = [im.copy() for im in img_list]
        batches, idx_maps = self._preprocess_cls(img_list)
        outputs = []
        for batch in batches:
            out = self.cls_session.run(None, {self.cls_input_name: batch})
            outputs.append(out[0])
        cls_res = self._postprocess_cls(outputs, idx_maps, len(img_list), img_list)
        return img_list, cls_res, 0

    # ---- Recognition ----

    def _preprocess_rec(self, img_crop_list):
        num = len(img_crop_list)
        width_list = [im.shape[1] / float(im.shape[0]) for im in img_crop_list]
        indices = np.argsort(np.array(width_list))
        batches, idx_maps = [], []
        for beg in range(0, num, self.rec_batch_num):
            end = min(num, beg + self.rec_batch_num)
            imgC, imgH, imgW = self.rec_image_shape
            max_wh_ratio = imgW / imgH
            for ino in range(beg, end):
                h, w = img_crop_list[indices[ino]].shape[:2]
                max_wh_ratio = max(max_wh_ratio, w / h)
            # Use fixed width if set by ONNX
            if self._rec_fixed_w > 0:
                max_wh_ratio = self._rec_fixed_w / imgH
            norm_list, idx_list = [], []
            for ino in range(beg, end):
                orig_idx = indices[ino]
                norm = _resize_norm_img(img_crop_list[orig_idx], self.rec_image_shape, max_wh_ratio=max_wh_ratio)
                norm_list.append(np.expand_dims(norm, axis=0))
                idx_list.append(orig_idx)
            if norm_list:
                batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
            idx_maps.append(idx_list)
        return batches, idx_maps

    def _postprocess_rec(self, batch_outputs, idx_maps, total_num):
        results = [("", 0.0)] * total_num
        # Decode each batch separately (different T per batch)
        for preds_batch, idx_list in zip(batch_outputs, idx_maps):
            texts = self._rec_post(preds_batch)
            for i, orig_idx in enumerate(idx_list):
                results[orig_idx] = texts[i]
        return results

    def recognize(self, img_crop_list):
        if not img_crop_list:
            return []
        batches, idx_maps = self._preprocess_rec(img_crop_list)
        outputs = []
        for batch in batches:
            out = self.rec_session.run(None, {self.rec_input_name: batch})
            outputs.append(out[0])
        return self._postprocess_rec(outputs, idx_maps, len(img_crop_list))

    # ---- Full pipeline ----

    def predict_image(self, image_path, visualize=False):
        img = cv2.imread(image_path)
        if img is None:
            raise FileNotFoundError(f"Cannot read image: {image_path}")
        return self(img, visualize=visualize)

    def __call__(self, img: np.ndarray, visualize=False, use_cls=None):
        ori_im = img.copy()

        boxes = self.detect(img)
        # print(f"[PPOCRv6] Detected {len(boxes)}")
        if len(boxes) == 0:
            return [] if not visualize else ori_im

        boxes = _sorted_boxes(boxes)
        # print(f"[PPOCRv6] Detected _sorted_boxes {len(boxes)}")

        img_crop_list = []
        for i, box in enumerate(boxes):
            crop = _get_rotate_crop_image(ori_im, np.array(box, dtype=np.float32))
            img_crop_list.append(crop)

        # Direction classifier
        do_cls = self.use_angle_cls if use_cls is None else use_cls
        if do_cls and self.cls_session is not None:
            img_crop_list, cls_res, _ = self.classify(img_crop_list)

        rec_res = self.recognize(img_crop_list)

        results = []
        for box, (text, conf) in zip(boxes, rec_res):
            if conf >= self.drop_score:
                results.append({"text": text, "confidence": round(conf, 4), "box": box.tolist()})

        return (results, draw_ocr_result(ori_im, results)) if visualize else (results, None)


# ============================================================================
# 7. CLI
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="PP-OCRv6 ONNX Inference (standalone, no Paddle dependency)")
    parser.add_argument("--det_onnx", type=str, default="axmodel/ax650/det_npu1.axmodel")
    parser.add_argument("--rec_onnx", type=str, default="axmodel/ax650/rec_npu1.axmodel")
    parser.add_argument("--char_dict", type=str, default="onnx/rec_inference.yml")
    parser.add_argument("--image", required=True, help="Input image path")
    parser.add_argument("--use_gpu", action="store_true")
    parser.add_argument("--drop_score", type=float, default=0.5)
    parser.add_argument("--det_limit_side_len", type=int, default=960)
    parser.add_argument("--det_db_thresh", type=float, default=0.2)
    parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
    parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
    parser.add_argument("--rec_batch_num", type=int, default=1)
    parser.add_argument("--resize_mode", type=str, default="letterbox", choices=["letterbox", "stretch"])
    # Classifier
    parser.add_argument("--use_angle_cls", action="store_true", help="Enable direction classifier")
    parser.add_argument("--cls_onnx", type=str, default="axmodel/ax650/cls_npu1.axmodel", help="Classifer ONNX model path")
    parser.add_argument("--cls_thresh", type=float, default=0.9, help="Angie classifier confidence threshold")
    parser.add_argument("--cls_batch_num", type=int, default=1)
    # Output
    parser.add_argument("--visualize", action="store_true")
    parser.add_argument("--output", type=str, default=None)
    parser.add_argument("--json", type=str, default=None)

    args = parser.parse_args()

    char_dict_src = args.char_dict
    if not os.path.exists(char_dict_src) and ("," in char_dict_src or char_dict_src.startswith("[")):
        char_dict = [c.strip() for c in char_dict_src.strip("[]").split(",") if c.strip()]
    else:
        char_dict = char_dict_src

    ocr = PPOCRv6Onnx(
        det_onnx=args.det_onnx,
        rec_onnx=args.rec_onnx,
        char_dict=char_dict,
        det_limit_side_len=args.det_limit_side_len,
        det_db_thresh=args.det_db_thresh,
        det_db_box_thresh=args.det_db_box_thresh,
        det_db_unclip_ratio=args.det_db_unclip_ratio,
        rec_batch_num=args.rec_batch_num,
        use_angle_cls=args.use_angle_cls,
        cls_onnx=args.cls_onnx,
        cls_thresh=args.cls_thresh,
        cls_batch_num=args.cls_batch_num,
        drop_score=args.drop_score,
        use_gpu=args.use_gpu,
        resize_mode=args.resize_mode,
    )

    do_viz = args.visualize or args.output is not None
    img = cv2.imread(args.image)
    if img is None:
        raise FileNotFoundError(f"Cannot read image: {args.image}")

    if do_viz:
        results, vis = ocr(img, visualize=True)
        out_path = args.output or "res-ax.jpg"
        cv2.imwrite(out_path, vis)
        print(f"Annotated image saved to: {out_path}")
    else:
        results, vis = ocr(img)

    if args.json:
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"Results saved to: {args.json}")
    else:
        for i, res in enumerate(results):
            print(f"{i+1}. {res['text']}  ({res['confidence']:.3f})")


if __name__ == "__main__":
    main()