import sys
import json
import os
import piexif
import sqlite3
from sqlite3 import Error
from PIL import Image
from deepface import DeepFace
from retinaface import RetinaFace
import numpy as np
import cv2

class NpEncoder(json.JSONEncoder):
  def default(self, obj):
    if isinstance(obj, np.integer):
      return int(obj)
    if isinstance(obj, np.floating):
      return float(obj)
    if isinstance(obj, np.ndarray):
      return obj.tolist()
        
models = ["VGG-Face", "Facenet", "Facenet512", "OpenFace", "DeepFace", "DeepID", "ArcFace", "Dlib", "SFace"]

model = DeepFace.build_model('ArcFace')
input_shape = DeepFace.functions.find_input_shape(model)


# Adapted from DeepFace
# https://github.com/serengil/deepface/blob/master/deepface/commons/functions.py
#
# Modified to use bicubic resampling and clip expansion, as well as to
# take a PIL Image instead of numpy array
def alignment_procedure(img, left_eye, right_eye):
  """
  Given left and right eye coordinates in image, rotate around point
  between eyes such that eyes are horizontal
  :param img: Image (not np.array)
  :param left_eye: Eye appearing on the left (right eye of person)
  :param right_eye: Eye appearing on the right (left eye of person)
  :return: adjusted image
  """
  dY = right_eye[1] - left_eye[1]
  dX = right_eye[0] - left_eye[0]
  radians = np.arctan2(dY, dX)
  rotation = 180 * radians / np.pi

  if True:
    img = img.rotate(
      angle = rotation,
      resample = Image.BICUBIC, 
      expand = True)

  return img

def extract_faces(img, threshold=0.75, model = None, allow_upscaling = True):
  faces = RetinaFace.detect_faces(
    img_path = img, 
    threshold = threshold, 
    model = model, 
    allow_upscaling = allow_upscaling)

  # Re-implementation of 'extract_faces' with the addition of keeping a 
  # copy of the face image for caching on disk
  if type(faces) == dict:
    k=1
    for key in faces:
      print(f'Processing face {k}/{len(faces)}')
      k+=1
      identity = faces[key]
      facial_area = identity["facial_area"]
      landmarks = identity["landmarks"]
      left_eye = landmarks["left_eye"]
      right_eye = landmarks["right_eye"]

      if False: # Draw the face rectangle and eyes
        cv2.rectangle(img, 
          (int(facial_area[0]), int(facial_area[1])),
          (int(facial_area[2]), int(facial_area[3])),
          (0, 0, 255), 2)      
        cv2.circle(img, (int(left_eye[0]), int(left_eye[1])), 5, (255, 0, 0), 2)
        cv2.circle(img, (int(right_eye[0]), int(right_eye[1])), 5, (0, 255, 0), 2)
    
      # Find center of face, then crop to square
      # of equal width and height
      width = facial_area[2] - facial_area[0]
      height = facial_area[3] - facial_area[1]
      x = facial_area[0] + width * 0.5
      y = facial_area[1] + height * 0.5

      # Make thumbnail a square crop
      if width > height:
        height = width
      else:
        width = height

      #width *= 1.25
      #height *= 1.25

      left = max(round(x - width * 0.5), 0)
      right = min(round(left + width), img.shape[1])  # Y is 1
      top = max(round(y - height * 0.5), 0)
      bottom = min(round(top + height), img.shape[0]) # X is 0

      left_eye[0] -= top
      left_eye[1] -= left
      right_eye[0] -= top
      right_eye[1] -= left

      facial_img = img[top: bottom, left: right]

      # Eye order is reversed as the routine does them backwards
      image = Image.fromarray(facial_img)
      image = alignment_procedure(image, right_eye, left_eye)
      image = image.resize(size = input_shape, resample = Image.LANCZOS)
      resized = np.asarray(image)
        
      identity['vector'] = DeepFace.represent(
        img_path = resized,
        model_name = 'ArcFace',
        model = model, # pre-built
        detector_backend = 'retinaface',
        enforce_detection = False)

      identity["face"] = { 
        'top': facial_area[1] / img.shape[0],
        'left': facial_area[0] / img.shape[1],
        'bottom': facial_area[3] / img.shape[0],
        'right': facial_area[2] / img.shape[1]
      }

      identity['image'] = Image.fromarray(resized)

    return faces

#face verification
#img_path = sys.argv[1]

def create_connection(db_file):
  """ create a database connection to the SQLite database
      specified by db_file
  :param db_file: database file
  :return: Connection object or None
  """
  conn = None
  try:
    conn = sqlite3.connect(db_file)
  except Error as e:
    print(e)

  return conn

def create_face(conn, face):
  """
  Create a new face in the faces table
  :param conn:
  :param face:
  :return: face id
  """
  sql = '''
  INSERT INTO faces(photoId,scanVersion,faceConfidence,top,left,bottom,right,descriptorId)
  VALUES(?,?,?,?,?,?,?,?)
  '''
  cur = conn.cursor()
  cur.execute(sql, (
    face['photoId'], 
    face['scanVersion'], 
    face['faceConfidence'],
    face['top'],
    face['left'],
    face['bottom'],
    face['right'],
    face['descriptorId']
  ))
  conn.commit()
  return cur.lastrowid

def create_face_descriptor(conn, face):
  """
  Create a new face in the faces table
  :param conn:
  :param face:
  :return: descriptor id
  """
  sql = '''
  INSERT INTO facedescriptors(descriptors)
  VALUES(?)
  '''
  cur = conn.cursor()
  cur.execute(sql, (np.array(face['vector']),))
  conn.commit()
  return cur.lastrowid

def update_face_count(conn, photoId, faces):
  """
  Update the number of faces that have been matched on a photo
  :param conn:
  :param photoId:
  :param faces:
  :return: None
  """
  sql = '''
  UPDATE photos SET faces=? WHERE id=?
  '''
  cur = conn.cursor()
  cur.execute(sql, (faces, photoId))
  conn.commit()
  return None

base = '/pictures/'
conn = create_connection('../db/photos.db')
with conn:
  cur = conn.cursor()
  res = cur.execute('''
    SELECT photos.id,photos.faces,albums.path,photos.filename FROM photos
    LEFT JOIN albums ON (albums.id=photos.albumId) 
    WHERE photos.faces=-1
    ''')
  rows = res.fetchall()
  count = len(rows)
  i=1
  for row in rows:
    photoId, photoFaces, albumPath, photoFilename = row
    img_path = f'{base}{albumPath}{photoFilename}'
    print(f'Processing {i}/{count}: {img_path}')
    i+=1
    img = Image.open(img_path)
    img = img.convert()
    img = np.asarray(img)
    faces = extract_faces(img)
    if faces is None:
      update_face_count(conn, photoId, 0)
      continue
    j=1
    for key in faces:
      face = faces[key]
      image = face['image']
      print(f'Writing face {j}/{len(faces)}')
      j+=1
      
      #face['analysis'] = DeepFace.analyze(img_path = img, actions = ['age', 'gender', 'race', 'emotion'], enforce_detection = False)
      #face['analysis'] = DeepFace.analyze(img, actions = ['emotion'])

      # TODO: Add additional meta-data allowing back referencing to original
      # photo
      face['version'] = 1 # version 1 doesn't add much...

      data = {k: face[k] for k in set(list(face.keys())) - set(['image', 'facial_area', 'landmarks'])}
      json_str = json.dumps(data, ensure_ascii=False, indent=2, cls=NpEncoder)

      faceDescriptorId = create_face_descriptor(conn, face)

      faceId = create_face(conn, {
        'photoId': photoId,
        'scanVersion': face['version'],
        'faceConfidence': face['score'],
        'top': face['face']['top'],
        'left': face['face']['left'],
        'bottom': face['face']['bottom'],
        'right': face['face']['right'],
        'descriptorId': faceDescriptorId,
      })

      path = f'faces/{faceId % 10}'
      try:
        os.mkdir(path)
      except FileExistsError:
        pass
        
      with open(f'{path}/{faceId}.json', 'w', encoding = 'utf-8') as f:
        f.write(json_str)

      # Encode this data into the JPG as Exif
      exif_ifd = {piexif.ExifIFD.UserComment: json_str.encode()}
      exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {},
        "thumbnail": None, "GPS": {}}
      image.save(f'{path}/{faceId}.jpg', exif = piexif.dump(exif_dict))

      #df = DeepFace.find(img, db_path = '/db')
      #print(df.head())

    update_face_count(conn, photoId, len(faces))