from time import sleep
from io import BytesIO
from picamera import PiCamera
import cv2
import numpy as np
import os
import sys
import json

picCount = 0
new = False


def start():
    print('Starting training..')
    label = int(input('Label (Integer): '))
    name = str(raw_input('Person name: '))
    picCount = int(input('Number of Photos: '))
    global new
    new = promptNew()
    if new:
        newFileName = str(raw_input('New File name: '))
        oldFileName = newFileName
    else:
        print('Found in faces:')
        print(os.listdir('faces/'))
        oldFileName = str(raw_input('Old File name: '))
        newFileName = str(raw_input('New File name: '))
    clearTrainFolder()
    dum = input('Press any key to start capture..')
    capture(picCount)
    addPerson(label, newFileName, oldFileName)
    addName(name)


def clearTrainFolder():
    print('clearing train folder')
    filelist = [f for f in os.listdir('train/') if f.endswith(".jpg")]
    for f in filelist:
        os.remove(os.path.join('train/', f))


def addPerson(label, newFileName, oldFileName):
    try:
        # faceFilePath = '/faces/hesham-saeed2'
        faceFile = 'faces/' + oldFileName
        face_recognizer = cv2.createLBPHFaceRecognizer()
        global new
        if new:
            print('new file. skipping load function')
        else:
            print('loading old file')
            face_recognizer.load(faceFile)
    except:
        print('[ Error ] Problem in init/loading LBPHfacerecognizer')
        sys.exit()
    images = []
    labels = []
    image_names = os.listdir('train/')
    for image_name in image_names:
        if image_name.startswith('.'):
            continue;
        image_path = 'train/' + image_name
        image = cv2.imread(image_path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)
        for (x, y, w, h) in faces:
            images.append(gray[y:y + h, x:x + w])
            labels.append(label)
            cv2.imshow('face', gray[y:y + h, x:x + w])
            cv2.waitKey(100)
            print('Learning ' + image_name)
    if new:
        face_recognizer.train(images, np.array(labels))
    else:
        face_recognizer.update(images, np.array(labels))
    face_recognizer.save('faces/' + newFileName)
    print('Updated and saved file in faces/' + newFileName)


def capture(count):
    try:
        print('Initializing camera')
        camera = PiCamera()
        camera.resolution = (640, 480)
        print('[ OK ] Camera')
        sleep(2)
    except:
        print('[ Error ] Can not initialize PiCamera')
        sys.exit()
    for i in range(1, count + 1):
        pic = 'train/' + str(i) + '.jpg'
        camera.capture(pic)
        photo = cv2.imread(pic, 1)
        cv2.imshow('Photo', photo)
        print('Captured ' + str(i) + '\nPath ' + pic)
        cv2.waitKey(100)
        sleep(1)
    cv2.destroyAllWindows()
    print('Done!')


def promptNew():
    filemode = raw_input('Make new File? (y/n): ')
    if filemode == 'y':
        return True
    elif filemode == 'n':
        return False
    else:
        print('incorrect input')
        promptNew()


def addName(newName):
    with open("faces/names.json", "r") as read_file:
        exist = False
        namesJson = json.load(read_file)
        names = list(namesJson)
        i = 0
        for name in names:
            if newName == name:
                exist = True
                print('name already in names list with index ' + str(i))
            i += 1
        if not exist:
            names.append(newName)
            print('Added ' + newName + 'to names.json file')
    with open("faces/names.json", "w") as write_file:
        json.dump(names, write_file)


start()