import os
import json
import cv2
import numpy as np
import mediapipe as mp
from rknnlite.api import RKNNLite  # 导入RKNNLite





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 = 1,  #如果想要检测双手，请设置成2
            min_detection_confidence = 0.5,  #手势关键点的阈值
            model_complexity = 0  #使用最简单的模型  如果效果不准确 可以考虑设置比较复制的模型  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):
        """加载RKNN模型并解析元数据"""
        # 创建RKNNLite实例
        self.rknn_lite = RKNNLite()
        
        # 加载RKNN模型
        ret = self.rknn_lite.load_rknn(model_path)
        if ret != 0:
            raise RuntimeError(f'加载RKNN模型失败, 错误码: {ret}')
        
        # 初始化运行时环境  强制使用npu  core_mask=RKNNLite.NPU_CORE_0
        ret = self.rknn_lite.init_runtime()
        if ret != 0:
            raise RuntimeError(f'初始化NPU运行时失败, 错误码: {ret}')
        
        # 从同目录的JSON文件加载元数据
        metadata_path = self._get_metadata_path(model_path)
        self._load_metadata(metadata_path)

    def _get_metadata_path(self, model_path):
        """获取元数据文件的绝对路径"""
        # 尝试与模型同目录的JSON文件
        base_dir = os.path.dirname(model_path)
        base_name = os.path.basename(model_path)
        metadata_name = os.path.splitext(base_name)[0] + 'rknn_metadata.json'
        metadata_path = os.path.join(base_dir, metadata_name)
        
        # 如果文件不存在，尝试默认名称
        if not os.path.exists(metadata_path):
            metadata_path = os.path.join(base_dir, 'rknn_metadata.json')
            
        return metadata_path
    
    def _load_metadata(self, metadata_path):
        """从JSON文件加载元数据"""
        try:
            with open(metadata_path, 'r', encoding='utf-8') as f:
                metadata = json.load(f)
                
            self.class_labels = metadata.get('classes', ["点赞", "点踩", "胜利", "拳头", "我爱你", "手掌"])
            min_max = metadata.get('minMax', {})
            self.min_vals = min_max.get('min', [])
            self.max_vals = min_max.get('max', [])
            
            print(f"从 {metadata_path} 加载元数据成功")
            print(f"类别标签: {self.class_labels}")
            
        except Exception as e:
            print(f"加载元数据失败: {e}")
            # 设置默认值
            self.class_labels = ["点赞", "点踩", "胜利", "拳头", "我爱你", "手掌"]
            self.min_vals = []
            self.max_vals = []

    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):
        """使用RKNN Lite进行手势分类推理
        Args:
            keypoints (np.ndarray): 输入数据, 形状为(63,)
        Returns:
            results (dict): 分类结果字典，包含'class_id', 'label', 'confidence'
        """
        # 预处理输入数据
        input_data = self.preprocess(keypoints)
        # 使用RKNN Lite进行推理
        outputs = self.rknn_lite.inference(inputs=[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()