import face_recognition
import glob
import drawing
import numpy as np
from PIL import Image, ImageDraw
from webcam_video_stream import WebcamVideoStream
import cv2

class Recognition:
    def __init__(self):

        # Variables setup
        self.known_face_encodings = []
        self.known_face_names = []
        self.current_face_encodings = None
        self.current_face_locations = None
        self.last_face_locations = None

        self.dirty = True

        # Loading known faces
        path = './img/known/'
        path_slice = slice(len(path),100)

        for filename in glob.glob(f"{path}*.jpg"):
            # Get the filename (without path/extension)
            name = filename[path_slice]
            extension_slice = slice(0, len(name)-4)
            name = name[extension_slice]

            # Load image in face_recognition
            current_image = face_recognition.load_image_file(filename)
            current_encoding = face_recognition.face_encodings(current_image)[0]

            # Add encodings and name to the arraylists
            self.known_face_encodings.append(current_encoding)
            self.known_face_names.append(name)

    def save_locations(self):
        if(self.current_face_locations != None):
            self.last_face_locations = self.current_face_locations

    def rescale_locations(self, img_scale):
        # Rescale the 1/4 img locations for the original image
        output_locations = []
        for location in self.current_face_locations:
            location = [x/img_scale for x in location]
            output_locations.append(location)
        self.current_face_locations = output_locations

    def read_locations(self, frame):
        self.current_face_locations = face_recognition.face_locations(frame)

        if(self.current_face_locations != None):
            print(len(self.current_face_locations), " face(s) found")

    def read_encodings(self, frame):
        if(self.dirty == True):
            print("dirty")
            self.current_face_encodings = face_recognition.face_encodings(frame, self.current_face_locations)
            self.dirty = False

    def label_faces(self, frame):
        pil_image = Image.fromarray(frame)

        draw = ImageDraw.Draw(pil_image)

        copy = self.current_face_encodings

        for(top, right, bottom, left), self.current_face_encodings in zip(self.current_face_locations, self.current_face_encodings):

            drawing.draw_on_frame(draw, top, right, bottom, left, self.current_face_encodings, self.current_face_locations, self.known_face_encodings, self.known_face_names)
        del draw

        opencv_image = np.array(pil_image)
        return opencv_image

    def dirtyflag_check(self):
        # If found more/less faces
        if((self.current_face_locations != None) & (self.last_face_locations != None)):
            if(len(self.current_face_locations) != len(self.last_face_locations)):
                dirty = True
                return
        else:
            return 0

        change_factor = 5
        loop_index = 0
        for location in self.current_face_locations:

            if((location[0] < self.last_face_locations[loop_index][0] + change_factor) and (location[0] > self.last_face_locations[loop_index][0] - change_factor)):
                #print("not enough changes")
                pass
            else:
                self.dirty = True
                #print("changes!")

            loop_index += 1


img_scale = 0.25

recognition = Recognition()

cap = WebcamVideoStream().start()

while(True):
    ret, frame = cap.read()

    newX,newY = frame.shape[1]*img_scale, frame.shape[0]*img_scale
    low_resolution_frame = cv2.resize(frame,(int(newX),int(newY)))
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    recognition.dirtyflag_check()

    recognition.read_locations(low_resolution_frame)

    recognition.read_encodings(low_resolution_frame)

    recognition.rescale_locations(img_scale)

    labeled_frame = recognition.label_faces(gray_frame)

    cv2.imshow('feed', labeled_frame)
    if cv2.waitKey(1) & 0xFF == 27:
        cap.stop()
        break

    recognition.save_locations()

cap.release()
cv2.destroyAllWindows()