#!/home/ubuntu/miniconda3/envs/act-env/bin/python
"""
人脸识别 ROS2 节点
订阅彩色相机话题，使用 DeepFace 进行人脸识别，发布人脸信息
支持向量数据库存储和动态添加人脸
"""
import sys
print("当前Python路径：", sys.executable)
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from std_msgs.msg import String
import cv2
import numpy as np
from deepface import DeepFace
import json
import os
import pickle
from pathlib import Path


class FaceRecognitionNode(Node):
    """人脸识别 ROS2 节点"""

    def __init__(self):
        super().__init__('face_recognition_node')
        
        # 订阅彩色相机话题
        self.subscription = self.create_subscription(
            Image,
            '/camera/cam_1/color/image_raw',
            self.image_callback,
            10
        )
        
        # 订阅 load_face 话题，用于接收要注册的人脸名称
        self.load_face_subscription = self.create_subscription(
            String,
            'load_face',
            self.load_face_callback,
            10
        )
        
        # 发布人脸信息话题
        self.publisher = self.create_publisher(
            String,
            'user_face',
            10
        )
        
        # 发布人脸位置话题
        self.face_location_publisher = self.create_publisher(
            String,
            'face_location',
            10
        )
        
        # 配置参数
        self.declare_parameter('detector_backend', 'opencv')
        self.declare_parameter('recognition_threshold', 0.4)
        self.declare_parameter('model_name', 'VGG-Face')
        
        # 获取参数
        self.detector_backend = self.get_parameter('detector_backend').value
        self.recognition_threshold = self.get_parameter('recognition_threshold').value
        self.model_name = self.get_parameter('model_name').value
        
        # 初始化人脸向量数据库
        self.face_database = {}
        self.face_embeddings = {}
        self.db_dir = Path('database')
        self.embedding_file = Path('face_embeddings.pkl')
        self.db_dir.mkdir(exist_ok=True)
        
        # 注册相关变量
        self._registration_name = None
        self._waiting_for_registration = False
        
        self.get_logger().info(f'人脸识别节点已初始化')
        self.get_logger().info(f'检测后端: {self.detector_backend}')
        self.get_logger().info(f'识别阈值: {self.recognition_threshold}')
        self.get_logger().info(f'模型名称: {self.model_name}')
        self.get_logger().info(f'数据库目录: {self.db_dir}')
        self.get_logger().info(f'向量存储文件: {self.embedding_file}')
        
        # 尝试加载向量数据库
        self.load_embeddings()
        # 加载图像数据库（用于注册新人脸时保存图像）
        self.load_database()
    
    def load_database(self):
        """加载人脸数据库"""
        if self.db_dir.exists():
            # 遍历数据库目录，加载所有人脸
            for person_dir in self.db_dir.iterdir():
                if person_dir.is_dir():
                    name = person_dir.name
                    image_paths = list(person_dir.glob('*.jpg')) + list(person_dir.glob('*.png'))
                    if image_paths:
                        self.face_database[name] = str(image_paths[0])
            
            self.get_logger().info(f'已加载 {len(self.face_database)} 个人脸')
    
    def load_embeddings(self):
        """加载人脸向量数据库"""
        if self.embedding_file.exists():
            try:
                with open(self.embedding_file, 'rb') as f:
                    self.face_embeddings = pickle.load(f)
                self.get_logger().info(f'从 {self.embedding_file} 加载了 {len(self.face_embeddings)} 个人脸向量')
            except Exception as e:
                self.get_logger().error(f'加载向量数据库失败: {e}')
                self.face_embeddings = {}
        else:
            self.get_logger().info(f'未找到向量数据库文件 {self.embedding_file}，将创建新数据库')
            self.face_embeddings = {}
    
    def save_embeddings(self):
        """保存人脸向量数据库"""
        try:
            with open(self.embedding_file, 'wb') as f:
                pickle.dump(self.face_embeddings, f)
            self.get_logger().info(f'人脸向量数据库已保存到 {self.embedding_file}')
            return True
        except Exception as e:
            self.get_logger().error(f'保存向量数据库失败: {e}')
            return False
    
    def extract_embedding(self, image):
        """提取人脸特征向量"""
        try:
            embedding = DeepFace.represent(
                img_path=image,
                model_name=self.model_name,
                detector_backend=self.detector_backend,
                enforce_detection=False
            )
            return np.array(embedding[0]['embedding'])
        except Exception as e:
            self.get_logger().error(f'提取特征向量时出错: {e}')
            return None
    
    def image_callback(self, msg):
        """处理图像回调"""
        try:
            # 将 ROS Image 消息转换为 OpenCV 图像
            image = self.ros_image_to_cv2(msg)
            
            # 检查是否需要注册人脸
            if self._waiting_for_registration and self._registration_name:
                self.register_face(image, self._registration_name)
                self._waiting_for_registration = False
                self._registration_name = None
                return
            
            # 分析人脸
            result = self.analyze_face(image)
            
            if result:
                # 发布人脸信息
                self.publish_face_info(result)
                
        except Exception as e:
            self.get_logger().error(f'处理图像时出错: {str(e)}')
    
    def ros_image_to_cv2(self, msg):
        """将 ROS Image 消息转换为 OpenCV 图像"""
        # 从 ROS Image 消息中提取图像数据
        height = msg.height
        width = msg.width
        channels = 3
        image_data = np.frombuffer(msg.data, dtype=np.uint8).reshape((height, width, channels))
        
        # 转换颜色空间（从 RGB 到 BGR）
        image = cv2.cvtColor(image_data, cv2.COLOR_RGB2BGR)
        
        return image
    
    def analyze_face(self, image):
        """分析人脸属性和识别"""
        try:
            # 分析人脸属性
            analysis_result = DeepFace.analyze(
                image,
                actions=['age', 'gender', 'emotion', 'race'],
                detector_backend=self.detector_backend,
                enforce_detection=False
            )
            
            # 识别人脸
            name, confidence = self.recognize_face(image)
            
            # 获取人脸位置信息
            face_location = None
            if 'region' in analysis_result[0]:
                region = analysis_result[0]['region']
                face_location = {
                    'x': region['x'],
                    'y': region['y'],
                }
            
            # 构建结果
            result = {
                '姓名': name,
                '置信度': float(confidence),
                '年龄': float(analysis_result[0]['age']),
                '性别': analysis_result[0]['dominant_gender'],
                '情绪': analysis_result[0]['dominant_emotion'],
            }
            
            # 发布人脸位置信息
            if face_location:
                self.publish_face_location(face_location)
            
            return result
            
        except Exception as e:
            self.get_logger().error(f'分析人脸时出错: {str(e)}')
            return None
    
    def recognize_face(self, image):
        """使用向量数据库识别人脸"""
        if not self.face_embeddings:
            return "Unknown", 0.0
        
        try:
            # 提取输入图像的特征向量
            input_embedding = self.extract_embedding(image)
            if input_embedding is None:
                return "Unknown", 0.0
            
            # 计算与数据库中所有人脸的相似度
            min_distance = float('inf')
            recognized_name = "Unknown"
            
            for name, embedding in self.face_embeddings.items():
                # 计算欧氏距离
                distance = np.linalg.norm(input_embedding - embedding)
                if distance < min_distance:
                    min_distance = distance
                    recognized_name = name
            
            # 如果距离小于阈值，则认为是匹配
            if min_distance < self.recognition_threshold:
                return recognized_name, min_distance
            
            return "Unknown", 0.0
            
        except Exception as e:
            self.get_logger().error(f'识别人脸时出错: {e}')
            return "Unknown", 0.0
    
    def load_face_callback(self, msg):
        """处理 load_face 话题回调，注册新人脸"""
        try:
            name = msg.data.strip()
            if not name:
                self.get_logger().warn('收到空的 load_face 消息')
                return
            
            self.get_logger().info(f'收到注册人脸请求，姓名: {name}')
            
            # 等待下一帧图像
            self.get_logger().info('等待捕获人脸图像...')
            
            # 临时变量，用于存储捕获的图像
            self._registration_name = name
            self._waiting_for_registration = True
            
        except Exception as e:
            self.get_logger().error(f'处理 load_face 消息时出错: {e}')
    
    def publish_face_info(self, face_info):
        """发布人脸信息"""
        # 将人脸信息转换为 JSON 字符串
        face_info_json = json.dumps(face_info, ensure_ascii=False)
        
        # 创建并发布消息
        msg = String()
        msg.data = face_info_json
        self.publisher.publish(msg)
        
        # 记录日志
        self.get_logger().info(f'发布人脸信息: {face_info_json}')
    
    def publish_face_location(self, face_location):
        """发布人脸位置信息"""
        # 将人脸位置信息转换为 JSON 字符串
        face_location_json = json.dumps(face_location, ensure_ascii=False)
        
        # 创建并发布消息
        msg = String()
        msg.data = face_location_json
        self.face_location_publisher.publish(msg)
        
        # 记录日志
        self.get_logger().info(f'发布人脸位置: {face_location_json}')
    
    def register_face(self, image, name):
        """注册人脸到向量数据库"""
        try:
            # 提取特征向量
            embedding = self.extract_embedding(image)
            if embedding is None:
                self.get_logger().error('无法提取人脸特征向量，注册失败')
                return False
            
            # 保存图像到数据库
            person_dir = self.db_dir / name
            person_dir.mkdir(exist_ok=True)
            
            # 生成唯一文件名
            import time
            timestamp = int(time.time())
            image_path = person_dir / f'{name}_{timestamp}.jpg'
            cv2.imwrite(str(image_path), image)
            
            # 更新向量数据库
            self.face_embeddings[name] = embedding
            self.face_database[name] = str(image_path)
            
            # 保存向量数据库
            if self.save_embeddings():
                self.get_logger().info(f'成功注册人脸: {name}')
                self.get_logger().info(f'图像保存到: {image_path}')
                self.get_logger().info(f'向量已添加到数据库')
                return True
            else:
                self.get_logger().error('保存向量数据库失败')
                return False
                
        except Exception as e:
            self.get_logger().error(f'注册人脸时出错: {e}')
            return False


def main(args=None):
    """主函数"""
    rclpy.init(args=args)
    
    try:
        face_recognition_node = FaceRecognitionNode()
        rclpy.spin(face_recognition_node)
    except KeyboardInterrupt:
        pass
    finally:
        if 'face_recognition_node' in locals():
            face_recognition_node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
