from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from scipy import misc
import cv2
import sys
import os
import argparse
import numpy as np
import mxnet as mx
import random
import cv2
import sklearn
from sklearn.decomposition import PCA
from time import sleep
from easydict import EasyDict as edict
from mtcnn_detector import MtcnnDetector
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src', 'common'))
import face_image
import face_preprocess

ga_modelPath = 'model/model'

ctx = mx.gpu(0)
threshold = 1.24
det_minsize = 50
det_threshold = [0.6,0.7,0.8]
mtcnn_path = os.path.join(os.path.dirname(__file__), 'mtcnn-model')
detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=det_threshold)

def get_model(ctx, image_size, model_str, layer):
    prefix = model_str
    epoch = 0
    sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
    all_layers = sym.get_internals()
    sym = all_layers[layer+'_output']
    model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
    model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
    model.set_params(arg_params, aux_params)
    return model


def get_ga(ga_model, face_img):
    ret = detector.detect_face(face_img, det_type = 0)
    if ret is None:
        return None

    bbox, points = ret

    if bbox.shape[0] == 0:
        return None
    bbox = bbox[0,0:4]
    points = points[0,:].reshape((2,5)).T
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    aligned = np.transpose(nimg, (2,0,1))

    input_blob = np.expand_dims(aligned, axis=0)
    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    ga_model.forward(db, is_train=False)
    ret = ga_model.get_outputs()[0].asnumpy()
    g = ret[:,0:2].flatten()
    gender = np.argmax(g)
    a = ret[:,2:202].reshape( (100,2) )
    a = np.argmax(a, axis=1)
    age = int(sum(a))

    return gender, age


def findAgeGender(img):
    image_size = img.shape
    print(image_size)
    ga_model = get_model(ctx, image_size, ga_modelPath, 'fc1')
    return get_ga(ga_model,img)

#print(findAgeGender(cv2.imread('Tom_Hanks_54745.png')))