'''
Uses models from: https://data.vision.ee.ethz.ch/cvl/rrothe/imdb-wiki/

@article{Rothe-IJCV-2016,
  author = {Rasmus Rothe and Radu Timofte and Luc Van Gool},
  title = {Deep expectation of real and apparent age from a single image without facial landmarks},
  journal = {International Journal of Computer Vision (IJCV)},
  year = {2016},
  month = {July},
}
'''

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

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



class _FaceAgeGender:
  '''Predicts age from face crop
  https://data.vision.ee.ethz.ch/cvl/rrothe/imdb-wiki/
  '''

  dnn_size = (224,224)
  dnn_mean = (104.0, 177.0, 123.0)  # ?
  # authors used imagenet mean
  #dnn_mean = [103.939, 116.779, 123.68] 
  ages = np.arange(0, 101).reshape(101, 1)
  padding = 0.4

  def __init__(self, fp_prototxt, fp_model):
    self.log = logger_utils.Logger.getLogger()
    self.net = cv.dnn.readNetFromCaffe(fp_prototxt, fp_model)

  def _preprocess(self, im, bbox_norm):
    # isolate face ROI, expand bbox by 40% according to authors
    #   https://data.vision.ee.ethz.ch/cvl/rrothe/imdb-wiki/
    dim = im.shape[:2][::-1]
    roi = bbox_norm.expand(self.padding).to_dim(dim).to_xyxy()
    im_face_crop = im[roi[1]:roi[3], roi[0]:roi[2]]  # isolate face roi

    # resize for blob
    im_resized = cv.resize(im_face_crop, self.dnn_size)
    blob = cv.dnn.blobFromImage(im_resized, 1.0, self.dnn_size, self.dnn_mean)
    return blob


class FaceGender(_FaceAgeGender):

  # use "apparent" age models
  fp_model = join(cfg.DIR_MODELS_CAFFE, 'gender', 'gender.caffemodel')
  fp_prototxt = join(cfg.DIR_MODELS_CAFFE, 'gender', 'gender.prototxt')

  def __init__(self):
    super().__init__(self.fp_prototxt, self.fp_model)

  def predict(self, im, bbox_norm):
    '''Predicts gender from face crop
    :param im: (numpy.ndarray) BGR image
    :param bbox_dim: (BBox) dimensioned
    :returns (dict) with scores for male and female
    '''
    im_blob = self._preprocess(im, bbox_norm)
    self.net.setInput(im_blob)
    preds = self.net.forward()[0]
    return {'f': preds[0], 'm': preds[1]}


class FaceAgeApparent(_FaceAgeGender):

  # use "apparent" age models
  fp_model = join(cfg.DIR_MODELS_CAFFE, 'age_apparent', 'dex_chalearn_iccv2015.caffemodel')
  fp_prototxt = join(cfg.DIR_MODELS_CAFFE, 'age_apparent', 'age.prototxt')

  def __init__(self):
    super().__init__(self.fp_prototxt, self.fp_model)

  def predict(self, im, bbox_norm):
    '''Predicts apparent age from face crop
    :param im: (numpy.ndarray) BGR image
    :param bbox_dim: (BBox) dimensioned
    :returns (float) predicted age
    '''
    im_blob = self._preprocess(im, bbox_norm)
    self.net.setInput(im_blob)
    preds = self.net.forward()[0]
    age = preds.dot(self.ages).flatten()[0]
    return age


class FaceAgeReal(_FaceAgeGender):

  # use "real" age models
  fp_model = join(cfg.DIR_MODELS_CAFFE, 'age_real', 'dex_imdb_wiki.caffemodel')
  fp_prototxt = join(cfg.DIR_MODELS_CAFFE, 'age_real', 'age.prototxt')

  def __init__(self):
    super().__init__(self.fp_prototxt, self.fp_model)

  def predict(self, im, bbox_dim):
    '''Predicts apparent age from face crop
    :param im: (numpy.ndarray) BGR image
    :param bbox_dim: (BBox) dimensioned
    :returns (float) predicted age
    '''
    im_blob = self._preprocess(im, bbox_dim)
    self.net.setInput(im_blob)
    preds = self.net.forward()[0]
    age = preds.dot(self.ages).flatten()[0]
    return age