import cv2
import numpy as np
import onnxruntime as ort


class DensityEstimatorONNX:
    def __init__(self, model_path, input_size=None, device="cpu", providers=None):
        self.model_path = model_path
        self.input_size = input_size

        if providers is None:
            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("[Density] providers:", active_providers)
        if (
            device.lower() in ["cuda", "gpu"]
            and "CUDAExecutionProvider" not in active_providers
        ):
            raise RuntimeError(
                "请求使用 CUDA 推理，但密度模型会话未启用 "
                "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 "
                "onnxruntime-gpu。"
            )

        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        self.output_names = [o.name for o in self.session.get_outputs()]

        print("[Density] input_name:", self.input_name)
        print("[Density] input_shape:", self.input_shape)
        print("[Density] output_names:", self.output_names)

        if self.input_size is None:
            h = self.input_shape[2]
            w = self.input_shape[3]
            if isinstance(h, int) and isinstance(w, int):
                self.input_size = (w, h)
            else:
                self.input_size = (400, 400)

    def preprocess(self, image_bgr):
        original_h, original_w = image_bgr.shape[:2]

        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(image_rgb, self.input_size)

        img = resized.astype(np.float32) / 255.0

        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        img = (img - mean) / std

        img = img.transpose(2, 0, 1)
        img = np.expand_dims(img, axis=0).astype(np.float32)

        return img, original_w, original_h

    def infer(self, image_bgr):
        input_tensor, original_w, original_h = self.preprocess(image_bgr)

        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
        density = outputs[0]

        density = np.squeeze(density).astype(np.float32)

        if density.ndim != 2:
            density = density.reshape(density.shape[-2], density.shape[-1])

        count = float(density.sum())

        density_map = cv2.resize(density, (original_w, original_h))

        return {
            "count": count,
            "density_map": density_map
        }

    @staticmethod
    def colorize_density_map(density_map):
        d = density_map.copy()
        d = d - d.min()
        if d.max() > 0:
            d = d / d.max()
        d = (d * 255).astype(np.uint8)
        heatmap = cv2.applyColorMap(d, cv2.COLORMAP_JET)
        return heatmap