import cv2 as cv
import numpy as np
import os
import tkinter
from PIL import Image, ImageTk


def traindata(dir):
    count = 1

    listFiles = os.listdir("data_train/" + str(dir))

    path = "data_train/" + str(dir) + "/%d.jpg"

    img = cv.imread(path % count)

    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

    sift = cv.xfeatures2d.SIFT_create()

    kp, des = sift.detectAndCompute(gray, None)

    while (count < len(listFiles)):
        count = count + 1

        img1 = cv.imread(path % count)

        gray1 = cv.cvtColor(img1, cv.COLOR_BGR2GRAY)

        kp1, des1 = sift.detectAndCompute(gray1, None)

        des = np.concatenate((des, des1))

    return des


# method detect face and save face
def detectFace(image, face):
    try:
        img = cv.imread(image)
        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

        face_cascade = cv.CascadeClassifier("haarcascade_frontalface_default.xml")
        faces = face_cascade.detectMultiScale(gray)

        for (x, y, w, h) in faces:
            faceCrop = img[y:y + h, x:x + w]
            croppedImg = cv.resize(faceCrop, (60, 60))  # resize to 60x60 px

        cv.imwrite(face, croppedImg)
        return 1
    except:
        print('Detect face error')
        return 0


# method predict face
def predictFace():
    # get des, key image input
    img = "input.jpg"
    face = "face.jpg"
    detectFace(img, face)
    imgtest = cv.imread("face.jpg")

    graytest = cv.cvtColor(imgtest, cv.COLOR_BGR2GRAY)

    sift = cv.xfeatures2d.SIFT_create()

    kptest, destest = sift.detectAndCompute(graytest, None)

    # get train data
    listDirs = os.listdir("data_train")

    for i in range(0, len(listDirs)):
        des = traindata(i)

        if i == 0:
            trainData = des
            ketqua = np.full((len(des), 1), i)
        else:
            trainData = np.concatenate((trainData, des), axis=0).astype(np.float32)
            ketqua = np.concatenate((ketqua, np.full((len(des), 1), i)), axis=0).astype(np.float32)

    knn = cv.ml.KNearest_create()
    knn.train(trainData, 0, ketqua)
    temp, result, nearest, distance = knn.findNearest(destest, 1)

    # print("Result: {}\n".format(result))
    # print("Nearest: {}\n".format(nearest))
    # print("Distance: {}\n".format(distance))

    # predict
    counts = np.bincount(np.array(result).reshape(-1).astype(int))
    # print(np.argmax(counts))
    # print("Predict image's information")
    f = open('people.txt', 'r')
    people = f.readlines()
    return people[np.argmax(counts)].strip()


def handleButtonPredict():
    lb.configure(text=predictFace())
    return


# Camera
def handleButtonCamera():
    key = cv.waitKey(1)
    webcam = cv.VideoCapture(0)
    webcam.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
    webcam.set(cv.CAP_PROP_FRAME_HEIGHT, 720)
    while True:
        try:
            check, frame = webcam.read()
            cv.imshow("Capturing", frame)
            key = cv.waitKey(1)
            if key == ord('s'):
                cv.imwrite(filename='input.jpg', img=frame)
                webcam.release()
                cv.waitKey(1650)
                cv.destroyAllWindows()
                print("Image saved!")
                loadImage()
                break
            elif key == ord('q'):
                print("Turning off camera.")
                webcam.release()
                print("Camera off.")
                print("Program ended.")
                loadImage()
                cv.destroyAllWindows()
                break

        except(KeyboardInterrupt):
            print("Turning off camera.")
            webcam.release()
            print("Camera off.")
            print("Program ended.")
            loadImage()
            cv.destroyAllWindows()
            break


# Crawl
def handleButtonCrawl():
    # add person in file people.txt
    f = open('people.txt', 'a')
    f.write(txtName.get())
    f.close()

    # crawl
    listFolder = len(os.listdir("data_train/"))
    pathFolder = "data_train/" + str(listFolder)
    fileFace = 1
    if not os.path.exists(pathFolder):
        os.mkdir(pathFolder)
        # capture picture to crawl data face
        # key = cv.waitKey(1)
        webcam = cv.VideoCapture(0)
        webcam.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
        webcam.set(cv.CAP_PROP_FRAME_HEIGHT, 720)
        while True:
            try:
                check, frame = webcam.read()
                cv.imshow("Capturing", frame)
                key = cv.waitKey(1)
                if key == ord('s'):
                    if os.path.exists("image.jpg"):
                        os.remove("image.jpg")
                    cv.imwrite(filename='image.jpg', img=frame)
                    # detect face
                    image = "image.jpg"
                    face = pathFolder + "/" + str(fileFace) + ".jpg"
                    if (detectFace(image, face) == 1):
                        print("Face saved!")
                        fileFace += 1
                elif key == ord('q'):
                    if os.path.exists("image.jpg"):
                        os.remove("image.jpg")
                    print("Turning off camera.")
                    webcam.release()
                    print("Camera off.")
                    print("Program ended.")
                    cv.destroyAllWindows()
                    break

            except(KeyboardInterrupt):
                if os.path.exists("image.jpg"):
                    os.remove("image.jpg")
                print("Turning off camera.")
                webcam.release()
                print("Camera off.")
                print("Program ended.")
                cv.destroyAllWindows()
                break


# Load image
def loadImage():
    load = Image.open("input.jpg")
    load = load.resize((600, 400))
    render = ImageTk.PhotoImage(load)
    img = tkinter.Label(image=render)
    img.image = render
    img.place(x=100, y=0)


# UI
window = tkinter.Tk()
window.title("Face recognition")
window.geometry("800x600")

loadImage()

# button
btnPredict = tkinter.Button(window, text="Predict", command=handleButtonPredict)
btnPredict.place(x=300, y=500)
btnCamera = tkinter.Button(window, text="Camera", command=handleButtonCamera)
btnCamera.place(x=200, y=500)
btnCrawl = tkinter.Button(window, text="Crawl", command=handleButtonCrawl)
btnCrawl.place(x=110, y=500)

# label
lb = tkinter.Label(window)
lb.place(x=400, y=505)
lbName = tkinter.Label(window)
lbName.place(x=110, y=450)
lbName.configure(text="Name : ")

# text box
txtName = tkinter.Entry(window, width=20)
txtName.place(x=200, y=450)

window.mainloop()
