import cv2
import dlib
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.applications import ResNet50
from tensorflow.python.keras.layers import Dense

import config


def get_age_model():
    age_model = ResNet50(
        include_top=False,
        weights='imagenet',
        input_shape=(config.RESNET50_DEFAULT_IMG_WIDTH, config.RESNET50_DEFAULT_IMG_WIDTH, 3),
        pooling='avg'
    )

    prediction = Dense(units=101,
                       kernel_initializer='he_normal',
                       use_bias=False,
                       activation='softmax',
                       name='pred_age')(age_model.output)

    age_model = Model(inputs=age_model.input, outputs=prediction)
    return age_model


def get_model(ignore_age_weights=False):
    base_model = get_age_model()
    if not ignore_age_weights:
        base_model.load_weights(config.AGE_TRAINED_WEIGHTS_FILE)
        print('Loaded weights from age classifier')
    last_hidden_layer = base_model.get_layer(index=-2)

    base_model = Model(
        inputs=base_model.input,
        outputs=last_hidden_layer.output)
    prediction = Dense(1, kernel_initializer='normal')(base_model.output)

    model = Model(inputs=base_model.input, outputs=prediction)
    return model

def get_trained_model():
    weights_file = config.MODEL_WEIGHTS_PATH
    _model = get_model(ignore_age_weights=True)
    _model.load_weights(weights_file)
    return _model

model = get_trained_model()
#detector = dlib.get_frontal_face_detector()

graph = tf.get_default_graph()

def predict_bmi(detected_faces, img):
    input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_h, img_w, _ = np.shape(input_img)

    faces = np.empty((len(detected_faces), config.RESNET50_DEFAULT_IMG_WIDTH,
                      config.RESNET50_DEFAULT_IMG_WIDTH, 3))


    for i, d in enumerate(detected_faces):
        x1, y1, x2, y2, w, h = d.left(), d.top(), d.right() + 1, \
                               d.bottom() + 1, d.width(), d.height()
        xw1 = max(int(x1 - config.MARGIN * w), 0)
        yw1 = max(int(y1 - config.MARGIN * h), 0)
        xw2 = min(int(x2 + config.MARGIN * w), img_w - 1)
        yw2 = min(int(y2 + config.MARGIN * h), img_h - 1)

        faces[i, :, :, :] = cv2.resize(img[yw1:yw2 + 1, xw1:xw2 + 1, :], (
            config.RESNET50_DEFAULT_IMG_WIDTH,
            config.RESNET50_DEFAULT_IMG_WIDTH)) / 255.00

        with graph.as_default():
            predictions = model.predict(faces)
            return predictions