#!/usr/bin/env python3
"""
PP-OCRv6 Recognition AX Inference (one-by-one, batch=1 only)

Dependencies:
    numpy, opencv-python, axengine, pyyaml, tqdm

Usage:
    # Single image inference
    python ppocrv6_rec_ax-one.py --rec_onnx rec.axmodel --char_dict inference.yml --image crop.png

    # Batch evaluation
    python ppocrv6_rec_ax-one.py --rec_onnx rec.axmodel --char_dict inference.yml \\
        --label_file labels.txt --dataset_root ./crops/ \\
        --output_json result.json --verbose
"""

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

import cv2
import numpy as np
import axengine as ax
import yaml
from tqdm import tqdm


# ============================================================================
# 1. Utilities
# ============================================================================

def _edit_distance(pred: str, target: str) -> Tuple[int, float]:
    m, n = len(pred), len(target)
    if m == 0:
        return n, 1.0
    if n == 0:
        return m, 1.0
    dp = list(range(n + 1))
    for i in range(1, m + 1):
        prev = dp[0]
        dp[0] = i
        for j in range(1, n + 1):
            temp = dp[j]
            if pred[i - 1] == target[j - 1]:
                dp[j] = prev
            else:
                dp[j] = 1 + min(prev, dp[j], dp[j - 1])
            prev = temp
    distance = dp[n]
    normalized = distance / max(m, n)
    return distance, normalized


def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
    if isinstance(source, list):
        return list(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."
        )


def _resize_norm_img(
    img: np.ndarray,
    image_shape: Tuple[int, int, int] = (3, 48, 320),
) -> np.ndarray:
    imgC, imgH, imgW = image_shape
    h, w = img.shape[:2]
    max_wh_ratio = imgW * 1.0 / imgH
    ratio = w * 1.0 / h
    max_wh_ratio = max(max_wh_ratio, ratio)
    max_wh_ratio = min(max_wh_ratio, imgW / imgH)
    target_w = int(imgH * max_wh_ratio)

    if math.ceil(imgH * ratio) > target_w:
        resized_w = target_w
    else:
        resized_w = int(math.ceil(imgH * ratio))

    resized = cv2.resize(img, (resized_w, imgH))
    resized = resized.astype("float32")
    resized = resized.transpose((2, 0, 1)) 
    # resized /= 255.0
    # resized -= 0.5
    # resized /= 0.5
    padded = np.zeros((imgC, imgH, target_w), dtype=np.float32)
    padded[:, :, 0:resized_w] = resized
    return padded


# ============================================================================
# 2. CTC Decoder
# ============================================================================

class _CTCLabelDecode:
    """CTC greedy decoder for recognition output."""

    def __init__(self, character_list: List[str], use_space_char: bool = 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
        self.dict = {char: i for i, char in enumerate(dict_character)}

    def decode(
        self,
        text_index: np.ndarray,
        text_prob: Optional[np.ndarray] = None,
        is_remove_duplicate: bool = True,
    ) -> List[Tuple[str, float]]:
        result_list = []
        batch_size = len(text_index)
        for batch_idx in range(batch_size):
            selection = np.ones(len(text_index[batch_idx]), dtype=bool)
            if is_remove_duplicate:
                selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
            selection &= text_index[batch_idx] != 0
            char_list = [
                self.character[int(tid)]
                for tid in text_index[batch_idx][selection]
            ]
            if text_prob is not None:
                conf_list = text_prob[batch_idx][selection]
            else:
                conf_list = np.ones(len(selection), dtype=np.float32)
            if len(conf_list) == 0:
                conf_list = np.array([0.0], dtype=np.float32)
            text = "".join(char_list)
            result_list.append((text, float(np.mean(conf_list))))
        return result_list

    def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
        preds_idx = preds.argmax(axis=2)
        preds_prob = preds.max(axis=2)
        return self.decode(preds_idx, preds_prob, is_remove_duplicate=True)


# ============================================================================
# 3. Recognition Engine (one-by-one)
# ============================================================================

class PPOCRv6RecOne:
    """Recognition engine for batch=1 AX model.
    
    Flow: load image → preprocess single → inference → decode.
    """

    def __init__(
        self,
        rec_model: str,
        char_dict: Union[str, List[str]],
        rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
    ):
        self.rec_image_shape = rec_image_shape
        self.session = ax.InferenceSession(rec_model)
        self.input_name = self.session.get_inputs()[0].name
        char_list = _load_char_dict(char_dict)
        self._decoder = _CTCLabelDecode(char_list, use_space_char=True)

    def _preprocess_one(self, img: np.ndarray) -> np.ndarray:
        tensor = _resize_norm_img(img, self.rec_image_shape)
        return np.expand_dims(tensor, axis=0).astype(np.float32)

    def infer_one(self, img: np.ndarray) -> Tuple[str, float]:
        inp = self._preprocess_one(img)
        out = self.session.run(None, {self.input_name: inp})[0]
        results = self._decoder(out)
        return results[0] if results else ("", 0.0)

    def __call__(
        self, images: List[np.ndarray]
    ) -> List[Tuple[str, float]]:
        results = []
        for img in tqdm(images, desc="Inference", unit="sample"):
            results.append(self.infer_one(img))
        return results

    def predict_image(self, path: str) -> Tuple[str, float]:
        im = cv2.imread(path)
        if im is None:
            raise FileNotFoundError(f"Cannot read: {path}")
        return self.infer_one(im)


# ============================================================================
# 4. Evaluation (one-by-one)
# ============================================================================

def evaluate(
    ocr: PPOCRv6RecOne,
    label_file: str,
    dataset_root: str = "",
    ignore_space: bool = True,
    verbose: bool = False,
) -> dict:
    images = []
    targets = []

    with open(label_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 2:
                continue
            img_path = os.path.join(dataset_root, parts[0].strip())
            images.append(img_path)
            targets.append(parts[1].strip())

    total = len(images)
    if total == 0:
        print("[WARN] No samples found in label file.")
        return {"acc": 0.0, "norm_edit_dis": 0.0, "total": 0, "correct": 0, "per_sample": []}

    correct = 0
    total_edit_dis = 0.0
    per_sample = []

    for i in tqdm(range(total), desc="Evaluating", unit="sample"):
        p = images[i]
        gt = targets[i]
        im = cv2.imread(p)
        if im is None:
            print(f"[WARN] Cannot read {p}, skipping.")
            pred, conf = ("", 0.0)
        else:
            pred, conf = ocr.infer_one(im)

        pred_clean = pred.replace(" ", "") if ignore_space else pred
        gt_clean = gt.replace(" ", "") if ignore_space else gt
        dist, norm_dist = _edit_distance(pred_clean, gt_clean)
        total_edit_dis += norm_dist
        is_correct = pred_clean == gt_clean
        if is_correct:
            correct += 1

        sample = {
            "image": p,
            "pred": pred,
            "gt": gt,
            "confidence": round(conf, 4),
            "correct": is_correct,
            "edit_distance": int(dist),
            "norm_edit_dis": round(norm_dist, 4),
        }
        per_sample.append(sample)

        if verbose:
            status = " OK" if is_correct else "MIS"
            print(
                f"[{status}] pred={pred!r:<30} gt={gt!r:<30} "
                f"conf={conf:.4f} edit={int(dist)} ndis={norm_dist:.4f}"
            )

    acc = correct / total
    norm_edit_dis = 1.0 - total_edit_dis / total

    return {
        "acc": round(acc, 6),
        "norm_edit_dis": round(norm_edit_dis, 6),
        "total": total,
        "correct": correct,
        "per_sample": per_sample,
    }


# ============================================================================
# 5. CLI
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="PP-OCRv6 Recognition AX (one-by-one) – inference & evaluation"
    )
    parser.add_argument(
        "--rec_onnx", type=str,
        default="axmodel/ax650/rec_npu2.axmodel",
        help="Path to recognition AX model",
    )
    parser.add_argument(
        "--char_dict", type=str,
        default="onnx/rec_inference.yml",
        help="Character dictionary: .yml (PostProcess.character_dict), .txt, or comma-list",
    )
    parser.add_argument("--rec_image_shape", type=str, default="3,48,320",
                        help="Recognition input shape C,H,W (comma separated)")

    parser.add_argument("--image", type=str, default=None, help="Single crop image path")

    parser.add_argument("--label_file", type=str,
                        default='dataset/ocr_rec_dataset_examples/val.txt',
                        help="Label file (image_path<TAB>gt_text per line)")
    parser.add_argument("--dataset_root", type=str,
                        default="dataset/ocr_rec_dataset_examples",
                        help="Prefix directory for image paths in label file")

    parser.add_argument("--ignore_space", action="store_true", default=True,
                        help="Ignore spaces when comparing (default: True)")
    parser.add_argument("--verbose", action="store_true", help="Print per-sample results")
    parser.add_argument("--output_json", type=str, default=None,
                        help="Save results to JSON file")

    args = parser.parse_args()

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

    image_shape = tuple(int(v) for v in args.rec_image_shape.split(","))
    if len(image_shape) != 3:
        raise ValueError("--rec_image_shape requires 3 comma-separated integers")

    ocr = PPOCRv6RecOne(
        rec_model=args.rec_onnx,
        char_dict=char_dict,
        rec_image_shape=image_shape,
    )

    if args.image:
        text, conf = ocr.predict_image(args.image)
        print(f"text={text!r}  confidence={conf:.4f}")
        if args.output_json:
            import json
            with open(args.output_json, "w", encoding="utf-8") as f:
                json.dump({"text": text, "confidence": conf}, f, ensure_ascii=False, indent=2)
        return

    if args.label_file:
        metrics = evaluate(
            ocr,
            args.label_file,
            dataset_root=args.dataset_root,
            ignore_space=args.ignore_space,
            verbose=args.verbose,
        )

        print()
        print("=" * 60)
        print("Evaluation Results")
        print("=" * 60)
        print(f"  Total samples:         {metrics['total']}")
        print(f"  Correct (exact match): {metrics['correct']}")
        print(f"  Accuracy:              {metrics['acc']:.4f} ({metrics['acc']*100:.2f}%)")
        print(f"  Norm Edit Distance:    {metrics['norm_edit_dis']:.4f}")
        print("=" * 60)

        if args.output_json:
            import json
            out = {k: v for k, v in metrics.items() if k != "per_sample"}
            out["per_sample"] = metrics["per_sample"]
            with open(args.output_json, "w", encoding="utf-8") as f:
                json.dump(out, f, ensure_ascii=False, indent=2)
            print(f"\nResults saved to: {args.output_json}")
        return

    parser.error("Either --image or --label_file must be provided.")


if __name__ == "__main__":
    main()