"""
This is an example of using the k-nearest-neighbors (KNN) algorithm for face recognition.

When should I use this example?
This example is useful when you wish to recognize a large set of known people,
and make a prediction for an unknown person in a feasible computation time.

Algorithm Description:
The knn classifier is first trained on a set of labeled (known) faces and can then predict the person
in an unknown image by finding the k most similar faces (images with closet face-features under eucledian distance)
in its training set, and performing a majority vote (possibly weighted) on their label.

For example, if k=3, and the three closest face images to the given image in the training set are one image of Biden
and two images of Obama, The result would be 'Obama'.

* This implementation uses a weighted vote, such that the votes of closer-neighbors are weighted more heavily.

Usage:

1. Prepare a set of images of the known people you want to recognize. Organize the images in a single directory
   with a sub-directory for each known person.

2. Then, call the 'train' function with the appropriate parameters. Make sure to pass in the 'model_save_path' if you
   want to save the model to disk so you can re-use the model without having to re-train it.

3. Call 'predict' and pass in your trained model to recognize the people in an unknown image.

NOTE: This example requires scikit-learn to be installed! You can install it with pip:

$ pip3 install scikit-learn

"""
import math
import os
import os.path
import pickle
from multiprocessing import Pool

import face_recognition
import numpy as np
from PIL import Image, ImageDraw
from face_recognition.face_recognition_cli import image_files_in_folder
from sklearn import neighbors

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


def listdir_nohidden(path):
    for f in os.listdir(path):
        if not f.startswith('.'):
            yield f


def train_worker(img_path, verbose, class_dir):
    while True:
        image = face_recognition.load_image_file(img_path)
        face_bounding_boxes = face_recognition.face_locations(image)
        encoding = []
        if len(face_bounding_boxes) != 1:
            # If there are no people (or too many people) in a training image, skip the image.
            if verbose:
                print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(
                    face_bounding_boxes) < 1 else "Found more than one face"))
            return [-1, -1]
        else:
            # Add face encoding for current image to the training set
            encoding = face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0]
            # print("Image {} has been trained".format(class_dir))
        return [class_dir, encoding]


def test_worker(image_path, model_path, label):
    predictions = predict(image_path, model_path=model_path)
    if predictions == []:
        # print("Test image No.{} cannot find label".format(label));
        return (label, [[-1]])
    # print("Looking for faces in {}".format(label))
    return label, predictions


def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):
    """
    Trains a k-nearest neighbors classifier for face recognition.

    :param train_dir: directory that contains a sub-directory for each known person, with its name.

     (View in source code to see train_dir example tree structure)

     Structure:
        <train_dir>/
        ├── <person1>/
        │   ├── <somename1>.jpeg
        │   ├── <somename2>.jpeg
        │   ├── ...
        ├── <person2>/
        │   ├── <somename1>.jpeg
        │   └── <somename2>.jpeg
        └── ...

    :param model_save_path: (optional) path to save model on disk
    :param n_neighbors: (optional) number of neighbors to weigh in classification. Chosen automatically if not specified
    :param knn_algo: (optional) underlying data structure to support knn.default is ball_tree
    :param verbose: verbosity of training
    :return: returns knn classifier that was trained on the given data.
    """

    def save_clf(distance):
        print("Training KNN classifier of metric:{}...".format(distance.__name__))
        knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, n_jobs=1, algorithm="auto",
                                                 weights="distance", metric=distance)
        knn_clf.fit(X, y)
        if model_save_path is not None:
            with open((model_save_path + distance.__name__ + ".clf"), 'wb') as f:
                pickle.dump(knn_clf, f)

    X = []
    y = []
    pool = Pool()
    tasks = list()
    if os.path.isfile((model_save_path + manhattan.__name__ + ".clf")) and os.path.isfile(
            (model_save_path + euclidean.__name__ + ".clf")) and os.path.isfile(
        (model_save_path + chebyshev.__name__ + ".clf")) and os.path.isfile(
        model_save_path + laplace_kernel.__name__ + ".clf") and os.path.isfile(
        model_save_path + poly_kernel.__name__ + ".clf") and os.path.isfile(
        model_save_path + rbf_kernel.__name__ + ".clf") and os.path.isfile(
        model_save_path + lin_kernel.__name__ + ".clf"):
        print('Use exist clf!')
        return
    # Loop through each person in the training set
    for class_dir in listdir_nohidden(train_dir):
        if not os.path.isdir(os.path.join(train_dir, class_dir)):
            continue

        # Loop through each training image for the current person
        for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)):
            tasks.append((img_path, verbose, class_dir))
            # image = face_recognition.load_image_file(img_path)
            # face_bounding_boxes = face_recognition.face_locations(image)
            #
            # if len(face_bounding_boxes) != 1:
            #     # If there are no people (or too many people) in a training image, skip the image.
            #     if verbose:
            #         print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(face_bounding_boxes) < 1 else "Found more than one face"))
            # else:
            #     # Add fac e encoding for current image to the training set
            #     X.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0])
            #     y.append(class_dir)
            #     print("Image {} has been trained".format(class_dir))

    # Determine how many neighbors to use for weighting in the KNN classifier
    if n_neighbors is None:
        n_neighbors = int(round(math.sqrt(len(X))))
        if verbose:
            print("Chose n_neighbors automatically:", n_neighbors)
    for (class_dir, encoding) in pool.starmap(train_worker, tasks):
        if (class_dir == -1):
            continue
        X.append(encoding)
        y.append(class_dir)
    save_clf(chebyshev)
    save_clf(euclidean)
    save_clf(manhattan)
    save_clf(poly_kernel)
    save_clf(laplace_kernel)
    save_clf(rbf_kernel)
    save_clf(lin_kernel)
    return


def chebyshev(X, Y):
    res = np.max(np.abs(X - Y))
    return res


def manhattan(X, Y):
    res = np.sum(np.abs(X - Y))
    return res


def euclidean(X, Y):
    res = np.sum((X - Y) ** 2) ** 0.5
    return res


def rbf_kernel(X, Y):
    method = rbf
    tmp = method(X, X) - 2 * method(X, Y) + method(Y, Y)
    return tmp ** 0.5


def laplace_kernel(X, Y):
    method = laplace
    tmp = method(X, X) - 2 * method(X, Y) + method(Y, Y)
    return tmp ** 0.5


def poly_kernel(X, Y):
    method = poly
    tmp = method(X, X) - 2 * method(X, Y) + method(Y, Y)
    return tmp ** 0.5


def lin_kernel(X, Y):
    method = lin
    tmp = method(X, X) - 2 * method(X, Y) + method(Y, Y)
    return tmp ** 0.5


def rbf(X, Y):
    tmp = np.linalg.norm(X - Y) ** 2
    return np.exp(-tmp / (2 * rbf_para ** 2))


def laplace(X, Y):
    tmp = np.linalg.norm(X - Y)
    return np.exp(-tmp / laplace_para)


def poly(X, Y):
    tmp = np.dot(X, Y)
    return tmp ** poly_para


def lin(X, Y):
    tmp = np.dot(X, Y)
    return tmp


# remove threshold
def predict(X_img_path, knn_clf=None, model_path=None, distance_threshold=0.6):
    """
    Recognizes faces in given image using a trained KNN classifier

    :param X_img_path: path to image to be recognized
    :param knn_clf: (optional) a knn classifier object. if not specified, model_save_path must be specified.
    :param model_path: (optional) path to a pickled knn classifier. if not specified, model_save_path must be knn_clf.
    :param distance_threshold: (optional) distance threshold for face classification. the larger it is, the more chance
           of mis-classifying an unknown person as a known one.
    :return: a list of names and face locations for the recognized faces in the image: [(name, bounding box), ...].
        For faces of unrecognized persons, the name 'unknown' will be returned.
    """
    if not os.path.isfile(X_img_path) or os.path.splitext(X_img_path)[1][1:] not in ALLOWED_EXTENSIONS:
        raise Exception("Invalid image path: {}".format(X_img_path))

    if knn_clf is None and model_path is None:
        raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")

    # Load a trained KNN model (if one was passed in)
    if knn_clf is None:
        with open(model_path, 'rb') as f:
            knn_clf = pickle.load(f)

    # Load image file and find face locations
    X_img = face_recognition.load_image_file(X_img_path)
    X_face_locations = face_recognition.face_locations(X_img)

    # If no faces are found in the image, return an empty result.
    if len(X_face_locations) == 0:
        return []

    # Find encodings for faces in the test iamge
    faces_encodings = face_recognition.face_encodings(X_img, known_face_locations=X_face_locations)

    # Use the KNN model to find the best matches for the test face
    # closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1)

    # are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(X_face_locations))]

    # are_matches = np.ones(1000)

    # Predict classes and remove classifications that aren't within the threshold
    return [(pred, loc) for pred, loc in
            zip(knn_clf.predict(faces_encodings), X_face_locations)]


def show_prediction_labels_on_image(img_path, predictions):
    """
    Shows the face recognition results visually.

    :param img_path: path to image to be recognized
    :param predictions: results of the predict function
    :return:
    """
    pil_image = Image.open(img_path).convert("RGB")
    draw = ImageDraw.Draw(pil_image)

    for name, (top, right, bottom, left) in predictions:
        # Draw a box around the face using the Pillow module
        draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))

        # There's a bug in Pillow where it blows up with non-UTF-8 text
        # when using the default bitmap font
        name = name.encode("UTF-8")

        # Draw a label with a name below the face
        text_width, text_height = draw.textsize(name)
        draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255))
        draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255, 255))

    # Remove the drawing library from memory as per the Pillow docs
    del draw

    # Display the resulting image
    pil_image.show()


def test(distance):
    tasks = list()
    for image_dir in listdir_nohidden(test_dir):
        label = image_dir
        labels.append(label)
        image_dir = os.path.join(test_dir, image_dir)
        for image_path in [os.path.join(image_dir, f) for f in listdir_nohidden(image_dir)]:
            tasks.append((image_path, ("trained_knn_model_" + distance.__name__ + ".clf"), label))
            # predictions = predict(image_path, model_path="trained_knn_model.clf")

        # Find all people in the image using a trained classifier model
        # Note: You can pass in either a classifier file name or a classifier model instance

        # Print results on the console
        #     for name, (top, right, bottom, left) in predictions:
        #         result_labels.append(name)
        #         test_result.append(name==label)
        # print("- Found {} at ({}, {})".format(name, left, top))
        # Display results overlaid on an image
    for (label, predictions) in pool.starmap(test_worker, tasks):
        # if predictions:
        test_result.append(label == predictions[0][0])
        result_labels.append(predictions[0][0])
        # else:
        #     test_result.append(0)
        #     result_labels.append(-1)

    accuracy = np.mean(test_result)
    print("Accuracy for distance:{} is {}".format(distance.__name__, accuracy))


if __name__ == "__main__":
    rbf_para = 1
    laplace_para = 0.20
    poly_para = 3
    test_dir = './data/TEST'
    train_dir = './data/TRAIN'
    train(train_dir, model_save_path="trained_knn_model_", n_neighbors=2, verbose=True)
    print("Training complete!")
    labels = list()
    test_result = []
    result_labels = []
    pool = Pool()

    # test(euclidean)
    # test(manhattan)
    # test(chebyshev)
    # test(laplace_kernel)
    test(poly_kernel)
    # test(rbf_kernel)
    # test(lin_kernel)
