import torch
from typing import Optional, List
import numpy
from insightface.app import FaceAnalysis
from insightface.app.common import Face
import numpy as np

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个


class FaceAnalyser:
    def __init__(self):
        self.device = self._select_device()
        self._analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
        self._analyser.prepare(ctx_id=0, det_size=(640, 640))

    def _select_device(self):
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'
        
    def get_many_faces(self,frame: np.ndarray) -> Optional[List[Face]]:
        try:
            return self._analyser.get(frame)
        except Exception as e:
            print(f"{e}")
            return None

    def get_one_face(self,frame: np.ndarray, position: int = 0) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            try:
                return many_faces[position]
            except IndexError:
                return many_faces[-1]
        return None


    def find_similar_face(self,frame: np.ndarray, reference_face: Face) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            for face in many_faces:
                if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
                    distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
                    if distance < SIMILAR_FACE_DISTANCE:
                        return face
        return None