#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
人脸识别示例：特征提取 + 相似度比对
"""

import cv2
import numpy as np
import sys
import os

# 添加父目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from face_feature_extractor import FaceFeatureExtractor


class SimpleFaceRecognizer:
    """简单的人脸识别器"""

    def __init__(self, threshold=0.7):
        """
        初始化识别器

        Args:
            threshold: 相似度阈值，大于此值认为是同一人
        """
        self.extractor = FaceFeatureExtractor()
        self.threshold = threshold
        self.database = {}  # 存储已注册的人脸特征

    def register(self, name, image):
        """
        注册人脸

        Args:
            name: 姓名
            image: 人脸图像

        Returns:
            bool: 是否注册成功
        """
        feature = self.extractor.extract_single_feature(image)

        if feature is not None:
            self.database[name] = feature
            print(f"✓ {name} 注册成功")
            return True
        else:
            print(f"✗ {name} 注册失败（质量检查未通过）")
            return False

    def recognize(self, image):
        """
        识别人脸

        Args:
            image: 待识别的人脸图像

        Returns:
            tuple: (姓名, 相似度) 或 (None, 0.0)
        """
        if not self.database:
            print("✗ 数据库为空，请先注册人脸")
            return None, 0.0

        # 提取待识别人脸特征
        unknown_feature = self.extractor.extract_single_feature(image)

        if unknown_feature is None:
            print("✗ 未检测到人脸或质量不合格")
            return None, 0.0

        # 计算与数据库中所有人脸的相似度
        best_match = None
        best_similarity = 0.0

        for name, known_feature in self.database.items():
            # 计算余弦相似度（特征已归一化，所以直接点乘）
            similarity = np.dot(unknown_feature, known_feature)

            if similarity > best_similarity:
                best_similarity = similarity
                best_match = name

        # 判断是否超过阈值
        if best_similarity >= self.threshold:
            return best_match, best_similarity
        else:
            return None, best_similarity

    def get_database_info(self):
        """获取数据库信息"""
        return {
            'total_persons': len(self.database),
            'persons': list(self.database.keys())
        }


def main():
    """主函数"""

    print("=" * 60)
    print("人脸识别演示")
    print("=" * 60)

    # 创建识别器
    recognizer = SimpleFaceRecognizer(threshold=0.7)

    # 示例1: 注册人脸
    print("\n步骤1: 注册人脸到数据库")
    print("-" * 60)

    # 这里需要替换为实际的图像路径
    registration_data = [
        ("袁中群", "1.jpg"),
        ("袁慧", "yh.jpg"),
        ("曾其乐", "zql.jpg"),
    ]

    for name, image_path in registration_data:
        if os.path.exists(image_path):
            image = cv2.imread(image_path)
            if image is not None:
                recognizer.register(name, image)
        else:
            print(f"⚠ 图像文件不存在: {image_path}")

    # 显示数据库信息
    db_info = recognizer.get_database_info()
    print(f"\n数据库状态:")
    print(f"  已注册人数: {db_info['total_persons']}")
    print(f"  人员列表: {', '.join(db_info['persons'])}")

    # 示例2: 人脸识别
    print("\n步骤2: 识别人脸")
    print("-" * 60)

    test_images = [
        "2.jpg",
        # "test_person2.jpg",
        # "unknown_person.jpg",
    ]

    for test_image_path in test_images:
        if not os.path.exists(test_image_path):
            print(f"⚠ 测试图像不存在: {test_image_path}")
            continue

        test_image = cv2.imread(test_image_path)
        if test_image is None:
            print(f"✗ 无法读取图像: {test_image_path}")
            continue

        print(f"\n识别图像: {test_image_path}")
        name, similarity = recognizer.recognize(test_image)

        if name:
            print(f"✓ 识别结果: {name} (相似度: {similarity:.4f})")
        else:
            print(f"✗ 未识别到已注册人员 (最高相似度: {similarity:.4f})")

    # 示例3: 批量识别
    print("\n步骤3: 批量处理")
    print("-" * 60)

    batch_dir = "./test_faces/"
    if os.path.exists(batch_dir):
        results = []
        for filename in os.listdir(batch_dir):
            if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
                image_path = os.path.join(batch_dir, filename)
                image = cv2.imread(image_path)
                if image is not None:
                    name, similarity = recognizer.recognize(image)
                    results.append({
                        'filename': filename,
                        'name': name,
                        'similarity': similarity
                    })

        # 输出批量结果
        print(f"批量处理完成，共处理 {len(results)} 张图像")
        for result in results:
            status = "✓" if result['name'] else "✗"
            name = result['name'] if result['name'] else "未知"
            print(f"{status} {result['filename']}: {name} ({result['similarity']:.4f})")
    else:
        print(f"⚠ 批量处理目录不存在: {batch_dir}")


if __name__ == "__main__":
    main()