from os.path import join
from pathlib import Path

import cv2 as cv
import numpy as np
import imutils

from app.utils import im_utils, logger_utils
from app.models.bbox import BBox
from app.settings import app_cfg as cfg
from app.settings import types
from app.models.bbox import BBox


# ----------------------------------------------------------------------
#
# 2D landmarks: 5pt and 68pt
#
# ----------------------------------------------------------------------

class Landmarks2D:

  # Abstract class

  def __init__(self):
    self.log = logger_utils.Logger.getLogger()

  def landmarks(self, im, bbox):
    # override
    self.log.warn('Define landmarks() function')
    pass

  def flatten(self, points):
    '''Converts list of point-tupes into a flattened list for CSV
    :param points: (list) of x,y points
    :returns dict item for each point (eg {'x1':100, 'y1':200})
    '''
    points_formatted = {}
    for idx, pt in enumerate(points, 1):
      for j, d in enumerate('xy'):
        points_formatted[f'{d}{idx}'] = pt[j]
    return points_formatted

  def normalize(self, points, dim):
    return [np.array(p)/dim for p in points]  # divides each point by w,h dim



import face_alignment

class FaceAlignment2D_68(Landmarks2D):

  # https://github.com/1adrianb/face-alignment
  # Estimates 2D facial landmarks
  
  def __init__(self, gpu=0, flip_input=False):
    t = face_alignment.LandmarksType._2D
    device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
    self.fa = face_alignment.FaceAlignment(t, device=device, flip_input=flip_input)
    super().__init__()
    self.log.debug(f'{device}')
    self.log.debug(f'{t}')

  def landmarks(self, im):
    '''Calculates the 2D facial landmarks
    :param im: (numpy.ndarray) BGR image
    :returns (list) of 68 (int) (tuples) as (x,y) 
    '''
    # predict landmarks
    points = self.fa.get_landmarks(im)  # returns array of arrays of 68 2D pts/face
    # convert to data type
    points = [list(map(int, p)) for p in points[0]]
    return points


class Dlib2D(Landmarks2D):
  
  def __init__(self, model):
    super().__init__()
    # init dlib
    import dlib
    self.predictor = dlib.shape_predictor(model)
    self.log.info(f'loaded predictor model: {model}')

  def landmarks(self, im, bbox):
    # Draw high-confidence faces
    dim_wh = im.shape[:2][::-1]
    bbox = bbox.to_dlib()
    im_gray = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
    points = [[p.x, p.y] for p in self.predictor(im_gray, bbox).parts()]
    return points


class Dlib2D_68(Dlib2D):
  
  def __init__(self):
    # Get 68-point landmarks using DLIB
    super().__init__(cfg.DIR_MODELS_DLIB_68PT)


class Dlib2D_5(Dlib2D):
  
  def __init__(self):
    # Get 5-point landmarks using DLIB
    super().__init__(cfg.DIR_MODELS_DLIB_5PT)


class MTCNN2D_5(Landmarks2D):
  
  # Get 5-point landmarks using MTCNN
  # https://github.com/ipazc/mtcnn
  # pip install mtcnn

  def __init__(self):
    super().__init__()
    self.log.warn('NB: MTCNN runs both face detector and landmark predictor together.')
    self.log.warn('   this will use face with most similar ROI')
    from mtcnn.mtcnn import MTCNN
    self.detector = MTCNN()

  def landmarks(self, im, bbox):
    '''Detects face using MTCNN and returns (list) of BBox
    :param im: (numpy.ndarray) image
    :returns list of BBox
    '''
    results = []
    dim_wh = im.shape[:2][::-1]  # (w, h)
    
    # run MTCNN to get bbox and landmarks
    dets = self.detector.detect_faces(im)
    keypoints = []
    bboxes = []
    #iterate detections and convert to BBox
    for det in dets:
      #rect = det['box']
      points = det['keypoints']
      # convert to normalized for contain-comparison
      points_norm = [np.array(pt)/dim_wh for pname, pt in points.items()]
      contains = False not in [bbox.contains(pn) for pn in points_norm]
      if contains:
        results.append(points)  # append original points
  
    return results


# ----------------------------------------------------------------------
#
# 3D landmarks
#
# ----------------------------------------------------------------------

class Landmarks3D:

  def __init__(self):
    self.log = logger_utils.Logger.getLogger()

  def landmarks(self, im, bbox):
    pass

  def flatten(self, points):
    '''Converts list of point-tupes into a flattened list for CSV
    :param points: (list) of x,y points
    :returns dict item for each point (eg {'x1':100, 'y1':200})
    '''
    points_formatted = {}
    for idx, pt in enumerate(points, 1):
      for j, d in enumerate('xyz'):
        points_formatted[f'{d}{idx}'] = pt[j]
    return points_formatted

  def normalize(self, points, dim):
    return [np.array(p)/dim for p in points]  # divides each point by w,h dim

    
class FaceAlignment3D_68(Landmarks3D):

  # Estimates 3D facial landmarks
  import face_alignment

  def __init__(self, gpu=0, flip_input=False):
    super().__init__()
    device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=flip_input)

  def landmarks(self, im, as_type=str):
    '''Calculates the 3D facial landmarks
    :param im: (numpy.ndarray) BGR image
    :returns (list) of 68 (int) (tuples) as (x,y, z) 
    '''
    # predict landmarks
    points = self.fa.get_landmarks(im)  # returns array of arrays of 68 3D pts/face
    # convert to data type
    points = [list(map(int, p)) for p in points[0]]
    return points