import os
import cv2
import numpy as np
import face_recognition
from pathlib import Path
import time

class FaceDetector:
    def __init__(self, conf_thres=0.5):
        """初始化人脸检测器，使用face_recognition库
        
        Args:
            conf_thres: 置信度阈值
        """
        self.conf_thres = conf_thres
        
    def _load_model(self):
        """不再需要加载模型，使用face_recognition库"""
        print("使用face_recognition库进行人脸检测，无需加载模型")
        pass
    
    def detect_faces(self, image):
        """使用face_recognition库检测图像中的人脸
        
        Args:
            image: 输入图像 (numpy数组，BGR格式)
            
        Returns:
            list: 检测到的人脸边界框列表 [x1, y1, x2, y2, confidence]
        """
        boxes = []
        
        # 使用face_recognition库检测人脸位置
        # 将BGR图像转换为RGB格式（face_recognition需要RGB格式）
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 使用face_recognition检测人脸位置
        face_locations = face_recognition.face_locations(rgb_image)
        
        # 将检测结果转换为所需的格式 [x1, y1, x2, y2, confidence]
        for face_location in face_locations:
            # face_location格式为(top, right, bottom, left)
            top, right, bottom, left = face_location
            
            # 转换为[x1, y1, x2, y2, confidence]格式
            # 由于face_recognition不提供置信度，我们使用1.0作为默认值
            boxes.append([left, top, right, bottom, 1.0])
        
        print(f"检测到 {len(boxes)} 个人脸")
        
        # 如果未检测到人脸，尝试使用OpenCV的人脸检测器作为备份
        if len(boxes) == 0:
            try:
                print("使用OpenCV备用人脸检测...")
                face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                faces = face_cascade.detectMultiScale(gray, 1.1, 5)
                
                for (x, y, w, h) in faces:
                    boxes.append([x, y, x+w, y+h, 0.8])  # 使用默认置信度0.8
                print(f"OpenCV检测到 {len(faces)} 个人脸")
            except Exception as e:
                print(f"OpenCV备用人脸检测失败: {e}")
        
        return boxes

class FaceRecognizer:
    def __init__(self, detector=None, similarity_threshold=0.6):
        """初始化人脸识别器
        
        Args:
            detector: 人脸检测器实例，如果为None则创建默认检测器
            similarity_threshold: 人脸相似度阈值
        """
        self.detector = detector if detector else FaceDetector()
        self.similarity_threshold = similarity_threshold
        # 使用OpenCV内置的人脸识别器
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.face_size = (128, 128)  # 标准化人脸大小
    
    def extract_face_features(self, image, box=None):
        """提取人脸特征
        
        Args:
            image: 输入图像
            box: 可选的人脸边界框 [x1, y1, x2, y2]，如果未提供则自动检测
            
        Returns:
            numpy.ndarray: 人脸特征向量
        """
        # 如果没有提供边界框，使用检测器检测人脸
        if box is None:
            boxes = self.detector.detect_faces(image)
            if not boxes:
                return None
            box = boxes[0][:4]  # 使用第一个检测到的人脸
        
        # 裁剪人脸区域
        x1, y1, x2, y2 = box
        face_img = image[y1:y2, x1:x2]
        
        # 转换为灰度图像
        face_gray = cv2.cvtColor(face_img, cv2.COLOR_BGR2GRAY)
        
        # 调整大小为标准尺寸
        face_resized = cv2.resize(face_gray, self.face_size)
        
        # 使用直方图均衡化增强对比度
        face_normalized = cv2.equalizeHist(face_resized)
        
        # 简单特征提取：将图像展平为向量
        # 在实际应用中，可以使用更复杂的特征提取方法，如HOG或深度特征
        feature_vector = face_normalized.flatten().astype(np.float32)
        
        # 对特征向量进行标准化
        if np.linalg.norm(feature_vector) > 0:
            feature_vector = feature_vector / np.linalg.norm(feature_vector)
        
        return feature_vector
    
    def compare_faces(self, known_feature, unknown_feature):
        """比较两个人脸特征的相似度
        
        Args:
            known_feature: 已知人脸特征
            unknown_feature: 待比对人脸特征
            
        Returns:
            float: 相似度 (0-1之间，越大表示越相似)
        """
        if known_feature is None or unknown_feature is None:
            return 0.0
        
        # 计算欧氏距离
        face_distance = np.linalg.norm(known_feature - unknown_feature)
        # 将距离转换为相似度
        similarity = 1.0 / (1.0 + face_distance)
        
        return similarity
    
    def identify_face(self, image, known_faces, box=None):
        """识别人脸
        
        Args:
            image: 输入图像
            known_faces: 已知人脸特征列表 [(person_id, name, feature_vector), ...]
            box: 可选的人脸边界框
            
        Returns:
            tuple or None: (person_id, name, similarity) 或 None（如果未匹配）
        """
        # 提取人脸特征
        face_feature = self.extract_face_features(image, box)
        if face_feature is None:
            return None
        
        # 寻找最佳匹配
        best_match = None
        best_similarity = 0
        
        for person_id, name, known_feature in known_faces:
            similarity = self.compare_faces(known_feature, face_feature)
            
            if similarity > best_similarity:
                best_similarity = similarity
                best_match = (person_id, name, similarity)
        
        # 如果最佳匹配的相似度低于阈值，返回None
        if best_match and best_match[2] >= self.similarity_threshold:
            return best_match
        
        return None

def draw_face_box(image, box, name=None, similarity=None, color=(0, 255, 0), thickness=2):
    """在图像上绘制人脸边界框和标签
    
    Args:
        image: 输入图像
        box: 人脸边界框 [x1, y1, x2, y2, ...]
        name: 可选的名称标签
        similarity: 可选的相似度标签
        color: 边界框颜色 (BGR格式)
        thickness: 边界框线条粗细
    
    Returns:
        numpy.ndarray: 绘制了边界框的图像
    """
    x1, y1, x2, y2 = box[:4]
    
    # 绘制边界框
    cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
    
    # 如果提供了名称，则绘制标签
    if name or similarity is not None:
        label = ""
        if name:
            label += name
        if similarity is not None:
            label += f" ({similarity:.2f})"
        
        # 设置标签背景
        (label_width, label_height), baseline = cv2.getTextSize(
            label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
        )
        cv2.rectangle(
            image,
            (x1, y1 - label_height - 5),
            (x1 + label_width, y1),
            color,
            -1  # 填充矩形
        )
        
        # 绘制标签文本
        cv2.putText(
            image,
            label,
            (x1, y1 - 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 0, 0),  # 黑色文本
            1
        )
    
    return image