import os
import json
import cv2
import numpy as np
import mediapipe as mp
import onnxruntime as ort




class Workflow:
    def __init__(self, model_path):
        # 确保model_path是绝对路径
        self.model_path = os.path.abspath(model_path)


        # 初始化MediaPipe Hands
        self.mp_hands = mp.solutions.hands
        self.hands = self.mp_hands.Hands(
            static_image_mode = False,  # 视频流模式  如果只是获取照片的手势关键点 请设置为True
            max_num_hands = 2,  #如果想要检测双手，请设置成2
            min_detection_confidence = 0.5,  #手势关键点的阈值
            model_complexity = 1  #使用最简单的模型  如果效果不准确 可以考虑设置比较复制的模型  1
        )
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles

        # 初始化元数据
        self.min_vals = None
        self.max_vals = None
        self.class_labels = None

        # 加载模型和元数据
        self.load_model(self.model_path)

    
    def load_model(self, model_path=None):
        """加载ONNX模型及元数据"""
        # 判断onnx是否存在
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"模型文件不存在: {self.model_path}")

        self.model_path = os.path.abspath(model_path)
        self.session = ort.InferenceSession(self.model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
        self._load_metadata()


    
    def _load_metadata(self):
        model_meta = self.session.get_modelmeta()
        # 兼容不同ONNX Runtime版本的元数据读取
        if hasattr(model_meta, 'custom_metadata_map'):
            meta = model_meta.custom_metadata_map
            min_max = json.loads(meta.get('minMaxValues', '{"min":[],"max":[]}'))
            self.min_vals, self.max_vals = min_max['min'], min_max['max']
            self.class_labels = json.loads(meta.get('classes', '[]'))
        else:
            self.min_vals = self.max_vals = self.class_labels = []
            for prop in model_meta.metadata_props:
                if prop.key == 'minMaxValues':
                    min_max = json.loads(prop.value)
                    self.min_vals, self.max_vals = min_max['min'], min_max['max']
                elif prop.key == 'classes':
                    self.class_labels = json.loads(prop.value)
        
        # 无元数据时的默认类别
        if not self.class_labels:
            self.class_labels = ["点赞", "点踩", "胜利", "拳头", "我爱你", "手掌"]
        
        # print("元数据:", self.min_vals, self.max_vals, self.class_labels)

    def det_hands(self, image):
        """检测手势关键点
        Args:
            image (np.ndarray): 输入图像, RGB格式
        Returns:
            keypoints (np.ndarray): 手势关键点数组, 形状为(63,)
            image (np.ndarray): 绘制关键点后的图像
        """

        # 手势关键点提取
        results = self.hands.process(image) #处理三通道彩色图
        if results.multi_hand_landmarks:  # 绘制手部关键点
            for hand_landmarks in results.multi_hand_landmarks:
                self.mp_drawing.draw_landmarks(
                    image,
                    hand_landmarks,
                    self.mp_hands.HAND_CONNECTIONS,
                    self.mp_drawing_styles.get_default_hand_landmarks_style(),
                    self.mp_drawing_styles.get_default_hand_connections_style())

            # 只使用检测到的第一只手
            landmarks = results.multi_hand_world_landmarks[0]
            
            # 提取关键点坐标
            keypoints = []
            for landmark in landmarks.landmark:
                keypoints.extend([landmark.x, landmark.y, landmark.z])
            keypoints = np.array(keypoints, dtype=np.float32)

            return keypoints, image
        else:
            return None, image

    def preprocess(self, keypoints):
        """预处理关键点数据，归一化
        Args:
            keypoints (np.ndarray): 原始关键点数组, 形状为(63,)
        Returns:
            normalized_kps (np.ndarray): 归一化后的关键点数组, 形状为(63,)
        """
        if not self.min_vals or not self.max_vals or len(self.min_vals) != len(keypoints):
            # 如果没有归一化参数或长度不匹配，返回原始数据
            return keypoints
            
        normalized = []
        for i, value in enumerate(keypoints):
            min_val = self.min_vals[i]
            max_val = self.max_vals[i]
            if max_val - min_val > 1e-6:  # 避免除以零
                normalized.append((value - min_val) / (max_val - min_val))
            else:
                normalized.append(0.0)
        # 转换为NumPy数组
        normalized_kps = np.array(normalized, dtype=np.float32)
        # 准备输入数据 (1, 63) 形状
        input_data = normalized_kps.reshape(1, -1).astype(np.float32)
        return input_data
    
    def postprocess(self, outputs):
        """后处理RKNN输出结果, 获取分类结果
        Args:
            outputs (list): RKNN推理输出列表
        Returns:
            class_id (int): 预测的类别ID
            confidence (float): 预测的置信度
        """
        predictions = outputs[0][0]
        # 获取预测结果
        class_id = np.argmax(predictions)
        confidence = float(predictions[class_id])
        # 获取类别标签
        label = self.class_labels[class_id] if class_id < len(self.class_labels) else f"未知类别 {class_id}"
        results = {
            'class_id': class_id,
            'label': label,
            'confidence': confidence,
        }
        return results
    
    def inference(self, keypoints):
        """使用onnx进行手势分类推理
        Args:
            keypoints (np.ndarray): 输入数据, 形状为(63,)
        Returns:
            results (dict): 分类结果字典，包含'class_id', 'label', 'confidence'
        """
        # 预处理输入数据
        input_data = self.preprocess(keypoints)
        # 进行推理
        outputs = self.session.run([self.output_name], {self.input_name: input_data})
        # 后处理输出结果
        results = self.postprocess(outputs)

        return results

    def release(self):
        """释放资源"""
        if hasattr(self, 'rknn_lite'):
            self.rknn_lite.release()
            print("NPU资源已释放")

    def __del__(self):
        """析构函数自动释放资源"""
        self.release()