import os
import dlib
import cv2
import sys
import numpy as np

try:
    import cPickle  # Python2.
except ImportError:
    import _pickle as cPickle  # Python3.


def enroll_face_dataset():
    pwd = sys.path[0]
    PREDICTOR_PATH = pwd + '/Facial_models/shape_predictor_68_face_landmarks.dat'
    FACE_RECOGNITION_MODEL_PATH = pwd + '/Facial_models/dlib_face_recognition_resnet_model_v1.dat'

    faceDetector = dlib.get_frontal_face_detector()
    shapePredictor = dlib.shape_predictor(PREDICTOR_PATH)
    faceRecognizer = dlib.face_recognition_model_v1(FACE_RECOGNITION_MODEL_PATH)

    faceDatasetFolder = pwd + '/Facial_images/face_rec/train/'

    subfolders = []
    for x in os.listdir(faceDatasetFolder):
        xpath = os.path.join(faceDatasetFolder, x)
        if os.path.isdir(xpath):
            subfolders.append(xpath)

    nameLabelMap = {}
    labels = []
    imagePaths = []
    for i, subfolder in enumerate(subfolders):
        for x in os.listdir(subfolder):
            xpath = os.path.join(subfolder, x)
            if x.endswith('jpg'):
                imagePaths.append(xpath)
                labels.append(i)
                nameLabelMap[xpath] = subfolder.split('/')[-1]

    index = {}
    i = 0
    faceDescriptors = None
    for imagePath in imagePaths:
        # print("processing: {}".format(imagePath))
        img = cv2.imread(imagePath)

        faces = faceDetector(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

        # print("{} Face(s) found".format(len(faces)))

        for k, face in enumerate(faces):

            shape = shapePredictor(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), face)

            landmarks = [(p.x, p.y) for p in shape.parts()]

            faceDescriptor = faceRecognizer.compute_face_descriptor(img, shape)

            faceDescriptorList = [x for x in faceDescriptor]
            faceDescriptorNdarray = np.asarray(faceDescriptorList, dtype=np.float64)
            faceDescriptorNdarray = faceDescriptorNdarray[np.newaxis, :]

            if faceDescriptors is None:
                faceDescriptors = faceDescriptorNdarray
            else:
                faceDescriptors = np.concatenate((faceDescriptors, faceDescriptorNdarray), axis=0)

            index[i] = nameLabelMap[imagePath]
            i += 1

    # Write descriors and index to disk
    np.save(pwd + '/Facial_models/descriptors.npy', faceDescriptors)
    with open(pwd + '/Facial_models/index.pkl', 'wb') as f:
        cPickle.dump(index, f)


if __name__ == "__main__":
    enroll_face_dataset()