import os
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis


class FaceRecognizer:
    """
    封装InsightFace人脸识别功能，支持从文件夹加载已知人脸。
    """

    def __init__(self, known_faces_dir: str):
        self.known_faces_dir = known_faces_dir
        self.app = self._initialize_insightface()
        self.known_faces_embeddings = {}
        self.known_faces_names = []
        self._load_known_faces()

    def _initialize_insightface(self):
        """初始化InsightFace FaceAnalysis应用"""
        print("初始化InsightFace引擎...")
        try:
            app = FaceAnalysis(name='buffalo_l', root='~/.insightface')
            app.prepare(ctx_id=0, det_size=(640, 640))
            print("InsightFace引擎初始化完成")
            return app
        except Exception as e:
            print(f"InsightFace初始化失败: {e}")
            print("请检查依赖是否安装及模型是否可访问")
            return None

    def _load_known_faces(self):
        """加载已知人脸特征"""
        if not os.path.exists(self.known_faces_dir):
            print(f"已知人脸目录不存在，已创建: {self.known_faces_dir}")
            os.makedirs(self.known_faces_dir, exist_ok=True)
            return

        print(f"从目录加载人脸特征: {self.known_faces_dir}")
        for person_name in os.listdir(self.known_faces_dir):
            person_dir = os.path.join(self.known_faces_dir, person_name)
            if os.path.isdir(person_dir):
                print(f"处理人物: {person_name}")
                embeddings = []
                for filename in os.listdir(person_dir):
                    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                        image_path = os.path.join(person_dir, filename)
                        try:
                            img = cv2.imread(image_path)
                            if img is None:
                                print(f"无法读取图片: {image_path}，已跳过")
                                continue

                            faces = self.app.get(img)
                            if faces:
                                embeddings.append(faces[0].embedding)
                                print(f"提取特征成功: {filename}")
                            else:
                                print(f"未检测到人脸: {filename}，已跳过")
                        except Exception as e:
                            print(f"处理图片出错 {image_path}: {e}")

                if embeddings:
                    self.known_faces_embeddings[person_name] = np.array(embeddings).mean(axis=0)
                    self.known_faces_names.append(person_name)
                    print(f"人物 {person_name} 加载完成，共 {len(embeddings)} 张照片")
                else:
                    print(f"人物 {person_name} 无有效特征，已跳过")
        print(f"人脸加载完成，共 {len(self.known_faces_names)} 人")

    def recognize(self, frame, threshold=0.4):
        """识别人脸并返回结果"""
        if not self.app or not self.known_faces_names:
            return False, None, None

        faces = self.app.get(frame)
        if not faces:
            return False, None, None

        for face in faces:
            for known_name in self.known_faces_names:
                known_embedding = self.known_faces_embeddings[known_name]

                embedding1 = face.embedding.astype(np.float32)
                embedding2 = known_embedding.astype(np.float32)

                dot_product = np.dot(embedding1, embedding2)
                norm_embedding1 = np.linalg.norm(embedding1)
                norm_embedding2 = np.linalg.norm(embedding2)

                similarity = 0.0 if (norm_embedding1 == 0 or norm_embedding2 == 0) else (
                        dot_product / (norm_embedding1 * norm_embedding2)
                )

                if similarity >= threshold:
                    print(f"检测到已知人物: {known_name} (相似度: {similarity:.4f})")
                    return True, known_name, similarity

        return False, None, None

    def test_single_image(self, image_path: str, threshold=0.4):
        """测试单张图片识别"""
        if not os.path.exists(image_path):
            print(f"图片不存在: {image_path}")
            return False, None, None

        frame = cv2.imread(image_path)
        if frame is None:
            print(f"无法读取图片: {image_path}")
            return False, None, None

        result, name, similarity = self.recognize(frame, threshold)

        if result:
            print(f"识别结果: {name} (相似度: {similarity:.4f})")

            faces = self.app.get(frame)
            for face in faces:
                bbox = face.bbox.astype(int)
                cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
                text = f"{name}: {similarity:.2f}"
                cv2.putText(frame, text, (bbox[0], bbox[1] - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

            cv2.imshow('识别结果', frame)
            print("按任意键关闭窗口...")
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        else:
            print("未识别到已知人脸")

        return result, name, similarity

#
# if __name__ == "__main__":
#     recognizer = FaceRecognizer(known_faces_dir="known_faces")
#     test_image_path = r"F:\OCR\RapidOCR-main\known_faces\B\14sino-qiu02-master1050.jpg"
#     recognizer.test_single_image(test_image_path, threshold=0.4)