import numpy as np
import cv2 as cv
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import os
import pandas as pd
import dlib
from joblib import dump, load

## Versions used for packages in this code
# sklearn V 1.1.0
# numpy V 
# OpenCV V 
# pandas V 
# pickle V 


# Directory Stuff
train_path = 'Training Images'

# Load model
detector = dlib.get_frontal_face_detector() 
sp = dlib.shape_predictor( 'shape_predictor_5_face_landmarks.dat' )
facerec = dlib.face_recognition_model_v1( 'dlib_face_recognition_resnet_model_v1.dat' )

# Encode array
encode_array = []

greg_path = 'Training Images\Greg'
josh_path = 'Training Images\Josh'
train_path = ' Training Images'
# Encode images
num_proc = 0
geg_proc = 0
for f in os.listdir( greg_path ):
  #print("Processing file: {}".format(f))
  path = os.path.join(greg_path, f)
  targ_img = dlib.load_rgb_image( path )
  t_dets = detector(targ_img, 1)
  if len( t_dets ) != 0 :
      shape = sp(targ_img, t_dets[0])
      num_proc += 1
      geg_proc += 1
      encode_array.append( np.array( facerec.compute_face_descriptor(targ_img, shape) ) ) # Append encoded image

for f in os.listdir( josh_path ):
  #print("Processing file: {}".format(f))
  path = os.path.join(josh_path, f)
  targ_img = dlib.load_rgb_image( path )
  t_dets = detector(targ_img, 1)
  if len( t_dets ) != 0 :
      shape = sp(targ_img, t_dets[0])
      num_proc += 1
      encode_array.append( np.array( facerec.compute_face_descriptor(targ_img, shape) ) ) # Append encoded image
      
df = pd.DataFrame( encode_array )
df.insert(df.shape[1], 'Geg', [1 if i < geg_proc else 0 for i in range(num_proc)] )
df.insert(df.shape[1], 'Josh', [0 if i < geg_proc else 1 for i in range(num_proc)] )

#encode_array = np.array( encode_array ) # Turn into numpy array 
#print(df[df.columns[df.shape[1]-3:df.shape[1]]])
df_shuf = df.sample(frac = 1)
#print( df_shuf ) 

# Separate train and validation sets 90:10 split
split = int(df_shuf.shape[0] * 0.9)
train = df_shuf.iloc[:split]
val = df_shuf.iloc[split:]

# Make train_x, train_y, val_x, val_y
labels = train.shape[1]-2
train_x = train[train.columns[:labels]]
train_y = train[train.columns[labels:]]

val_x = val[val.columns[:labels]]
val_y = val[val.columns[labels:]]

ground_truth = [ 1 if k==1 else 0 for k in val_y['Geg'] ]

# Load knn model
model = KNeighborsClassifier( n_neighbors = 3, algorithm = 'auto' )
print( train_x.to_numpy() )
model.fit(train_x.to_numpy(), train_y.to_numpy())

print(train_x.shape)
print(val_x.shape)
# Validate model
predict = model.predict( val_x )
print(predict)
predict_label = [ 1 if k[0]==1 else 0 for k in predict ]
print(predict_label)
# Score
score = accuracy_score( ground_truth, predict_label )
print( f'actual {ground_truth}\nPredicted {predict_label}' )
print(f" Accuracy score is {score}" )

# Save the model
dump( model, 'face_rec_test.joblib' )