# face_recognition_algorithm.py
from insightface.app import FaceAnalysis


class FaceRecognitionAlgorithm:
    """
    人脸识别核心算法类
    负责模型加载和设备配置，与业务逻辑分离
    """

    def __init__(self, model_name: str = 'buffalo_l', use_gpu: bool = True, use_npu: bool = False,
                 npu_device_id: int = 0):
        # 设备配置映射（NPU采用华为指定的完整参数）
        self.DEVICE_CONFIG = {
            "cpu": (['CPUExecutionProvider'], -1),
            "gpu": (['CUDAExecutionProvider'], 0),
            "npu": (
                [
                    (
                        "CANNExecutionProvider",
                        {
                            "device_id": npu_device_id,
                            "arena_extend_strategy": "kNextPowerOfTwo",
                            "npu_mem_limit": 16 * 1024 * 1024 * 1024,
                            "op_select_impl_mode": "high_precision",
                            "precision_mode": "allow_fp32_to_fp16",
                            "enable_cann_graph": True,
                        },
                    ),
                    "CPUExecutionProvider"
                ],
                npu_device_id
            )
        }

        # 算法参数设置
        self.det_size = 640  # 320快速 640中等 1280慢
        self.det_threshold = 0.5    # 人脸置信度

        # 根据设备类型选择配置
        if use_npu:
            device_type = "npu"
            print(f"✅ 使用NPU设备，设备ID: {npu_device_id}")
        elif use_gpu:
            device_type = "gpu"
            print("✅ 使用GPU设备")
        else:
            device_type = "cpu"
            print("✅ 使用CPU设备")

        providers, ctx_id = self.DEVICE_CONFIG[device_type]

        # 初始化人脸识别模型
        self.app = FaceAnalysis(name=model_name, providers=providers)
        self.app.prepare(
            ctx_id=ctx_id,
            det_thresh=self.det_threshold,
            det_size=(self.det_size, self.det_size)
        )

        print(f"✅ 人脸识别算法初始化完成 - 设备: {device_type.upper()}")

    def set_det_threshold(self, threshold: float):
        """设置检测阈值"""
        self.det_threshold = threshold
        print(f"✅ 检测阈值设置为: {threshold}")

    def set_det_size(self, size: int):
        """设置检测尺寸"""
        self.det_size = size
        print(f"✅ 检测尺寸设置为: {size}")

    def get_app(self) -> FaceAnalysis:
        """获取底层的FaceAnalysis实例"""
        return self.app