import face_recognition
from sklearn import svm
import joblib
import os
import concurrent.futures
import numpy as np
import time
import re

# Training the SVC classifier

# The training data would be all the face encodings from all the known images and the labels are their names
encodings = []
names = []
tuples = []
train_list = []
model_name = 'saved_model_2.pkl'
# known_people_folder = os.path.join('application_data', 'verification_images')
face_model = "large"
if face_model == "large":
    model_name = 'saved_model_2_large.pkl'

# dummy_data = [
#     {
#         "name": "Bayu",
#         "address": "299 St Louis Road Oak Forest, IL 60452",
#         "nik": "1000076456784631"
#     },
#     {
#         "name": "Dio",
#         "address": "22 Whitemarsh St. Mansfield, MA 02048",
#         "nik": "1000024792887549"
#     },
#     {
#         "name": "Hadi",
#         "address": "643 Honey Creek Dr. Milledgeville, GA 31061",
#         "nik": "1000038502830420"
#     },
#     {
#         "name": "Kevin",
#         "address": "881 Cooper Ave. Hummelstown, PA 17036",
#         "nik": "1000045356476664"
#     },
#     {
#         "name": "Matrix",
#         "address": "580 Glenwood Dr. Garner, NC 27529",
#         "nik": "1000023452134598"
#     },
#     {
#         "name": "Surya",
#         "address": "909 South St Paul Street Hopewell, VA 23860",
#         "nik": "1000075656784734"
#     },
# ]


def scan_known_people(known_people_folder):
    known_names = []
    known_face_encodings = []

    for file in image_files_in_folder(known_people_folder):
        img = face_recognition.load_image_file(file)
        encodings = face_recognition.face_encodings(img)

        if len(encodings) > 1:
            print("WARNING: More than one face found in {}. Only considering the first face.".format(file))

        if len(encodings) == 0:
            print("WARNING: No faces found in {}. Ignoring file.".format(file))
        else:
            known_face_encodings.append(encodings[0])

    return known_face_encodings


def image_files_in_folder(folder):
    img_list = [os.path.join(folder, f) for f in os.listdir(folder) if
                re.match(r'.*\.(jpg|jpeg|png|webp)', f, flags=re.I)]
    img_list.sort()
    return img_list


# known_faces = scan_known_people(known_people_folder)

# is_train = True
#
# if not is_train:
#     try:
#         clf = joblib.load(model_name)
#     except:
#         clf = None
#     if clf is None:
#         is_train = True


def train_image(image, person):
    face = face_recognition.load_image_file(image)
    face_bounding_boxes = face_recognition.face_locations(face)

    # If training image contains exactly one face
    if len(face_bounding_boxes) == 1:
        face_enc = face_recognition.face_encodings(face, model=face_model)[0]
        # Add face encoding for current image with corresponding label (name) to the training data
        tuples.append((face_enc, person))
    else:
        print(image + " was skipped and can't be used for training")

def train_svm(model_name):
    train_str = os.path.join('data', 'train')
    train_dir = os.listdir(os.path.join('data', 'train'))
    train_dir.sort()

    # for person in train_dir:
    #     pix_str = os.path.join(train_str, person)
    #     pix = os.listdir(os.path.join(train_str, person))
    #
    #     # Loop through each training image for the current person
    #     for person_img in pix:
    #         # Get the face encodings for the face in each image file
    #         image = os.path.join(pix_str, person_img)
    #         train_image(image, person)

    # Loop through each person in the training directory
    with concurrent.futures.ThreadPoolExecutor() as executor:
        for person in train_dir:
            pix_str = os.path.join(train_str, person)
            pix = os.listdir(os.path.join(train_str, person))

            # Loop through each training image for the current person
            for person_img in pix:
                # Get the face encodings for the face in each image file
                image = os.path.join(pix_str, person_img)
                executor.submit(train_image, image, person)

    # Create and train the SVC classifier

    encodings = [x for x, _ in tuples]
    names = [y for _, y in tuples]

    clf = svm.SVC(gamma='scale', probability=True)
    clf.fit(encodings, names)
    joblib.dump(clf, model_name)
    return 0

if __name__ == '__main__':
    train_svm(model_name)

# Load the test image with unknown faces into a numpy array
# test_image = face_recognition.load_image_file(os.path.join('application_data', 'input_image', 'input_image.jpg'))

# Find all the faces in the test image using the default HOG-based model
# face_locations = face_recognition.face_locations(test_image)
# no = len(face_locations)
# print("Number of faces detected: ", no)

# Predict all the faces in the test image using the trained classifier
# print("Found:")
# for i in range(no):
#     test_image_enc = face_recognition.face_encodings(test_image, model=face_model)[i]
#     start_time = time.perf_counter_ns()
#     proba_list = clf.predict_proba([test_image_enc])
#     dist = face_recognition.face_distance(known_faces, test_image_enc)
#     total = np.subtract(proba_list, dist)
#     end_time = time.perf_counter_ns()
#     process_time = end_time - start_time
#     classes = clf.classes_
#     print(classes)
#     print(total)
#     i = np.argmax(total)
#     proba = list(*proba_list)
#     name = dummy_data[i]["name"]
#     print(name, "{:.2f}".format(proba[i]))
#     print(proba)
#     print(list(dist))
#     print(list(*total))
#     print(process_time)