import face_recognition
import numpy as np
import cv2

from workers.mongo_worker import MongoWorker

class FaceRecognition():
    def __init__(self):
        self.mongo_worker = MongoWorker()
        self.face_encodings_in_cache = self.mongo_worker.get_all_encodings()

    def get_encodings(self, face_image):
        return face_recognition.face_encodings(face_image)

    def softmax(self, values):
        exp_values = np.exp(values)
        exp_values_sum = np.sum(exp_values)

        return exp_values / exp_values_sum

    def kl_divergence(self, face_encodings, face_to_compare):
        if len(face_encodings) == 0:
            return np.empty((0))

        face_encodings = np.asarray(self.softmax(face_encodings))
        face_to_compare = np.asarray(self.softmax(face_to_compare))

        face_encodings = face_encodings / face_encodings.sum(axis=1, keepdims=True)
        face_to_compare = face_to_compare / face_to_compare.sum()

        epsilon = 1e-10
        face_encodings = np.clip(face_encodings, epsilon, 1)
        face_to_compare = np.clip(face_to_compare, epsilon, 1)

        return np.sum(face_encodings * np.log(face_encodings / face_to_compare), axis=1)

    def face_distance(self, face_encodings:list[float]):
        result = {}
        if self.mongo_worker.count_persons() > len(self.face_encodings_in_cache):
            self.face_encodings_in_cache = self.mongo_worker.get_all_encodings()

        for person_id in self.face_encodings_in_cache:
            result.update({person_id: self.kl_divergence(face_encodings, self.face_encodings_in_cache[person_id])})

        print(result.values())
        try:
            minimal = min(result.values())
            print(minimal)
            if minimal > 0.001:
                return None
            return [key for key, val in result.items() if val == minimal][0]
        except:
            best_in_frame = dict()
            for key, elements in result.items():
                best_in_frame.update({key:min(elements)})
            minimal = min(best_in_frame.values())
            print(minimal)
            if minimal > 0.001:
                return None
            return [key for key, val in best_in_frame.items() if val == minimal][0]

if __name__ == "__main__":
    face_recognition_worker = FaceRecognition()
    img = cv2.imread("obama.jpg")
    print(face_recognition_worker.get_encodings(img))