import os
from os.path import join
from pathlib import Path

import cv2 as cv
import numpy as np
import dlib
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

def similarity(self, query_enc, known_enc):
    return np.linalg.norm(query_enc - known_enc, axis=1)

def flatten(vec):
    '''Converts N-D vector into a flattened list for CSV
    :param points: (list) a feature vector as list of floats
    :returns dict item for each point (eg {'d1':0.28442156, 'd1': 0.1868632})
    '''
    vec_flat = {}
    for idx, val in enumerate(vec, 1):
      vec_flat[f'd{idx}'] = val
    return vec_flat



class Extractor:

  n_dim = None  # override

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

  def flatten(self, vec):
    '''Converts N-D vector into a flattened list for CSV
    :param points: (list) a feature vector as list of floats
    :returns dict item for each point (eg {'d1':0.28442156, 'd1': 0.1868632})
    '''
    vec_flat = {}
    for idx, val in enumerate(vec, 1):
      vec_flat[f'd{idx}'] = val
    return vec_flat

  def unflatten_df(self, df):
    # convert from
    return [df[f'd{i}'] for i in range(1,257)]


class ExtractorVGG(Extractor):

  # https://github.com/ox-vgg/vgg_face2
  # Uses OpenCV DNN to extract feature vector for VGG Face 2 models
  n_dim = 256
  dnn_dim = (224,224)
  dnn_mean = (91.4953, 103.8827, 131.0912)

  def __init__(self):
    super().__init__()
    fp_model = '/data_store_hdd/apps/megapixels/models/caffe/vgg_face2/resnet50_256_caffe/resnet50_256.caffemodel'
    fp_prototxt = '/data_store_hdd/apps/megapixels/models/caffe/vgg_face2/resnet50_256_caffe/resnet50_256.prototxt'
    self.dnn = cv.dnn.readNetFromCaffe(fp_prototxt, fp_model)
    self.feat_layer = self.dnn.getLayerNames()[-2]

  def extract(self, im, bbox_norm, padding=0.3):
    '''Extracts feature vector for face crop
    :param im:
    :param bbox_norm: (BBox) normalized
    :param padding: (float) percent to extend ROI
    :param jitters: not used here
    :returns (list) of (float)'''
    
    bbox_ext = bbox_norm.expand(padding)
    dim = im.shape[:2][::-1]
    bbox_ext_dim = bbox_ext.to_dim(dim)
    x1,y1,x2,y2 = bbox_ext_dim.to_xyxy()
    im = im[y1:y2, x1:x2]
    # According to VGG, model trained using Bilinear interpolation (INTER_LINEAR)
    im = cv.resize(im, self.dnn_dim, interpolation=cv.INTER_LINEAR)
    blob = cv.dnn.blobFromImage(im, 1.0, self.dnn_dim, self.dnn_mean)
    self.dnn.setInput(blob)
    vec = np.array(self.dnn.forward(self.feat_layer)[0])
    vec_norm = np.array(vec)/np.linalg.norm(vec)  # normalize
    return vec_norm


class ExtractorDLIB(Extractor):

  # https://github.com/davisking/dlib/blob/master/python_examples/face_recognition.py
  # facerec.compute_face_descriptor(img, shape, 100, 0.25)
  # padding=opt_padding not yet implemented in dlib===19.16 but merged in master
  n_dim = 128
  process_width = 100

  def __init__(self, gpu=0, jitters=cfg.DLIB_FACEREC_JITTERS):
    super().__init__()
    self.num_jitters = cfg.DLIB_FACEREC_JITTERS
    # set and swap GPU visibility
    if gpu > -1:
      cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
      os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    self.predictor = dlib.shape_predictor(cfg.DIR_MODELS_DLIB_5PT)
    self.facerec = dlib.face_recognition_model_v1(cfg.DIR_MODELS_DLIB_FACEREC_RESNET)
    # unset and swap GPU visibility
    if gpu > -1:
      os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices  # reset GPU env


  def extract(self, im, bbox_norm):
    '''Converts image and bbox into 128d vector
    :param im: (numpy.ndarray) BGR image
    :param bbox_norm: (BBox) normalized
    '''
    # scale the image so the face is always 100x100 pixels
    dim = im.shape[:2][::-1]
    bbox_dim = bbox_norm.to_dim(dim)
    scale = self.process_width / bbox_dim.width
    cv.resize(im, None, fx=scale, fy=scale, interpolation=cv.INTER_LANCZOS4) 
    bbox_dim_dlib = bbox_dim.to_dlib()
    face_shape = self.predictor(im, bbox_dim_dlib)
    # this is only in dlib version 19.6++?
    # vec = self.facerec.compute_face_descriptor(im, face_shape, self.num_jitters, self.padding)
    # vectors are already normalized
    vec = self.facerec.compute_face_descriptor(im, face_shape, self.num_jitters)
    return vec