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


picCount = 0
new = False

def start():
	print('Starting training..')
	label = int(input('Label (Integer): '))
	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()
	print('Press any key to start capture..')
	dum = input('')
	capture(picCount)
	addPerson(label, newFileName, oldFileName)
	
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)
			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()
         
start()