from typing import List, Dict
import cv2
import numpy as np
import onnxruntime as ort


def softmax(x: np.ndarray) -> np.ndarray:
    x = x.astype(np.float32)
    x = x - np.max(x)
    exp = np.exp(x)
    return exp / (np.sum(exp) + 1e-12)


def to_probabilities(output: np.ndarray) -> np.ndarray:
    """
    兼容两种常见 ONNX 导出：
    1. 模型直接输出 logits，需要 softmax；
    2. 模型图中已经包含 softmax，直接输出概率。

    对概率再次做 softmax 会把原本很高的置信度压到接近均匀分布。
    """
    values = np.asarray(output, dtype=np.float32).reshape(-1)
    total = float(np.sum(values))
    is_probability_output = (
        values.size > 0
        and np.all(np.isfinite(values))
        and float(np.min(values)) >= -1e-5
        and float(np.max(values)) <= 1.0 + 1e-5
        and abs(total - 1.0) <= 1e-3
    )
    if is_probability_output:
        return values / max(total, 1e-12)
    return softmax(values)


class ConvNeXtONNXClassifier:
    def __init__(
        self,
        model_path: str,
        class_names: List[str],
        img_size: int = 224,
        device: str = "cpu",
    ):
        self.model_path = model_path
        self.class_names = class_names
        self.img_size = img_size

        providers = ["CPUExecutionProvider"]
        if device.lower() in ["cuda", "gpu"]:
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]

        self.session = ort.InferenceSession(model_path, providers=providers)
        active_providers = self.session.get_providers()
        print("[ConvNeXt] providers:", active_providers)
        if (
            device.lower() in ["cuda", "gpu"]
            and "CUDAExecutionProvider" not in active_providers
        ):
            raise RuntimeError(
                "请求使用 CUDA 推理，但 ConvNeXt 会话未启用 "
                "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 "
                "onnxruntime-gpu。"
            )
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]

        # 常见 ImageNet 归一化参数。若训练 ConvNeXt 时使用了其他 mean/std，需要保持一致。
        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)

    def preprocess(self, img_bgr: np.ndarray) -> np.ndarray:
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        img_rgb = cv2.resize(img_rgb, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
        img = img_rgb.astype(np.float32) / 255.0
        img = (img - self.mean) / self.std
        img = img.transpose(2, 0, 1)
        img = np.expand_dims(img, axis=0).astype(np.float32)
        return img

    def infer(self, img_bgr: np.ndarray) -> Dict:
        input_tensor = self.preprocess(img_bgr)
        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
        logits = np.asarray(outputs[0])

        if logits.ndim == 2:
            logits = logits[0]
        elif logits.ndim == 1:
            pass
        else:
            logits = logits.reshape(-1)

        probs = to_probabilities(logits)
        class_id = int(np.argmax(probs))
        confidence = float(probs[class_id])
        class_name = self.class_names[class_id] if 0 <= class_id < len(self.class_names) else str(class_id)

        return {
            "cls_class_id": class_id,
            "cls_class_name": class_name,
            "cls_confidence": confidence,
        }