import sys
import dlib
import cv2 as cv
import numpy as np
import os

''' target_path contains .png files of target faces 
	This will return a list of 128-D vector encodings for each target and corresponding list of names
'''
#def encode_targets( target_path, detector, shape_predictor, facerec, win ):
def encode_targets( target_path, detector, shape_predictor, facerec ):
  list = []
  names = []
  for face in os.listdir(target_path):
    img = cv.imread( os.path.join( target_path, face ) )
    # Get dimensions for image
    top = 0
    bot = img.shape[0]
    left = 0
    right = img.shape[1]
    # Make rect
    rect = dlib.rectangle( left, top, right, bot)
    print(f'Processing file { os.path.join( target_path, face ) }')
    shape = shape_predictor( img, rect )
    list.append(np.array( facerec.compute_face_descriptor( img, shape ) ) )
    names.append( face.split('.')[0] )
    
    # Add window to check if shape is right for encoding
    #img = dlib.load_rgb_image( os.path.join( target_path, face ) )
    #win.set_image( img )
    #win.add_overlay( shape )
    #dlib.hit_enter_to_continue()
    #win.clear_overlay()
  return names, list
	
  
if len(sys.argv) != 4:
    print(
        "Call this program like this:\n"
        "   ./face_alignment.py shape_predictor_5_face_landmarks.dat dlib_face_recognition_resnet_model_v1.dat Auth\n"
        "You can download a trained facial shape predictor from:\n"
        "    http://dlib.net/files/shape_predictor_5_face_landmarks.dat.bz2\n")
    print( len(sys.argv) )
    exit()
	
# Bounding Box variables
RED = ( 0, 0, 255 )
GREEN = ( 0, 255, 0 )
color = RED # Initialize color
downscale = 0.5
thresh = 0.6

predictor_path = sys.argv[1]
face_rec_model_path = sys.argv[2]

targets_path = sys.argv[3]

# Load models
detector = dlib.get_frontal_face_detector() # Face detector
sp = dlib.shape_predictor( predictor_path ) # Landmark detector (5-point)
facerec = dlib.face_recognition_model_v1( face_rec_model_path )


# Compute face encodings for target faces
#authorized_name, authorized_encoding = encode_targets( targets_path, detector, sp, facerec, win )
authorized_name, authorized_encoding = encode_targets( targets_path, detector, sp, facerec)

#print(authorized_name)
#print(authorized_encoding[0])

# Webcam setup
vid = cv.VideoCapture(0)

cap_frame = 0

bbs = []
names = []

while( True ):
  ret, frame = vid.read() # Capture video frame
  frame = cv.rotate( frame, cv.ROTATE_180) # Rotate frame
  
  # Facial recognition on 10th frame
  if cap_frame == 2 and ret:
    # Shrink frame to process
    frame_small = cv.resize( frame, (0, 0), fx= downscale, fy = downscale )
    #frame_gray = cv.cvtColor( frame_small, cv.COLOR_BGR2GRAY ) # Grayscale for dlib
    
    # Find bounding boxes of faces
    dets = detector( frame_small, 1 )
    
    # Print number of faces found
    num_faces = len(dets)
    #print(num_faces)
    names = []
    
    bbs = [ [  ( int(box.left()*(1 / downscale)) , int( box.top() * ( 1 / downscale ) ) ),  ( int( box.right() * (1 / downscale) ), int( box.bottom() * (1 / downscale) ) ) ] for box in dets ]
    # Process each face
    #shape = [ sp(frame_gray, box) for box in dets ]
    #unknown_encode = [ np.array( facerec.compute_face_descriptor( frame_gray, s ) ) for s in shape ]
    #distances = [ [ np.linalg.norm( ue - auth ) for auth in authorized_encoding ] for ue in unknown_encode] 
    
    print(num_faces)
    if( num_faces >= 1 ):
        print(bbs[0])
        print(bbs[0][0][0])
        saved_img = frame[ bbs[0][0][1]:bbs[0][1][1], bbs[0][0][0]:bbs[0][1][0] ]
        print(saved_img.shape)
        cv.imshow('frame', saved_img)
        cv.imwrite('Akugbe.png', saved_img)
        dlib.hit_enter_to_continue()
    #print( len(shape), len(unknown_encode), len(distances) )
    
    for box in dets:
      #left_top = ( box.left()*(1 / downscale), box.top() * ( 1 / downscale ) )
      #right_bot = ( box.right() * (1 / downscale), box.bottom() * (1 / downscale) )
      
      # Find shape and encode 
      shape = sp( frame_small, box )
      unknown_encode = np.array( facerec.compute_face_descriptor( frame_small, shape ) ) # Encode face
      name = 'unknown'
      # Determine if found face is authorized
      distances = [ np.linalg.norm( unknown_encode - auth ) for auth in authorized_encoding ] # Calc distance
      # Find smallest value
      #print(distances)
      min = np.argmin( distances )
      if distances[min] < thresh:
        name = authorized_name[min]
      names.append( name )
      
      
    
  # Draw bounding box
  #print(names)
  for i, box in enumerate(bbs):
    color = RED if names[i] == 'unknown' else GREEN 
    #print(color) 
    #print(names[i])
    cv.rectangle( frame, box[0], box[1], color , 2 )
  #cv.imshow('frame', frame)
  cv.imshow('frame', frame)
  key = cv.waitKey(1) 
  
  if key & 0xFF == ord('s'):
      saved_img = frame[ bbs[0][0][0]:bbs[0][0][1], bbs[0][1][0]:bbs[0][1][1] ]
      cv.imwrite('Akugbe.png', saved_img)
  cap_frame = 0 if cap_frame == 2 else cap_frame + 1
  if key & 0xFF == ord('q'):
    break
    

vid.release()
cv.destroyAllWindows()