import math
from sklearn import neighbors
import os
import os.path
import pickle
from PIL import Image, ImageDraw
import face_recognition
from face_recognition.face_recognition_cli import image_files_in_folder
import numpy as np

ALLOWED_EXTENSIONS = ('png', 'jpg', 'jpeg')


def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):
    encodings = []
    names = []

    train_dir = os.listdir(f"{train_dir}/")
    print(train_dir)

    for person in train_dir:
        pix = os.listdir(f"dataset/{person}")

        for person_img in pix:
            print(f"dataset/{person}/{person_img}")
            face = face_recognition.load_image_file(f"dataset/{person}/{person_img}")
            print(face.shape)

            height, width, _ = face.shape
            face_location = (0, width, height, 0)
            print(width, height)

            face_enc = face_recognition.face_encodings(face, known_face_locations=[face_location])

            face_enc = np.array(face_enc)
            face_enc = face_enc.flatten()

            encodings.append(face_enc)
            names.append(person)

    print(np.array(encodings).shape)
    print(np.array(names).shape)

    knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance')
    knn_clf.fit(encodings, names)

    if model_save_path is not None:
        with open(model_save_path, 'wb') as f:
            pickle.dump(knn_clf, f)

    return knn_clf


if __name__ == "__main__":
    print("Training KNN classifier...")
    classifier = train("dataset", model_save_path="trained_knn_model.clf", n_neighbors=2)
    print("Training Complete!")
