import gradio as gr
import tensorflow as tf
import cv2
import mtcnn
import numpy as np

model  = tf.keras.models.load_model('./model')

def load_and_preprocess_image(im_path, detector, maxWidth = 512):

    desiredLeftEye = (0.36, 0.43)

    # Load the image and convert it to grayscale
    try:
        image = cv2.imread(im_path)
    except:
        return 0

    if image is None:
        return 0

    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Detect the face in the image
    result = detector.detect_faces(image)

    # Get the bounding box for the face
    x, y, w, h = result[0]['box']

    desiredFaceWidth = 224
    desiredFaceHeight = 224
    
    # Get the landmarks for the face
    landmarks = result[0]['keypoints']
    
    # Calculate the angle between the eyes
    eye_1 = landmarks['left_eye']
    eye_2 = landmarks['right_eye']
    dy = eye_2[1] - eye_1[1]
    dx = eye_2[0] - eye_1[0]
    angle = np.arctan2(dy, dx) * 180 / np.pi

    desiredRightEyeX = 1.0 - desiredLeftEye[0]

    dist = np.sqrt((dx ** 2) + (dy ** 2))
    desiredDist = (desiredRightEyeX - desiredLeftEye[0]) * desiredFaceWidth
    scale = desiredDist / dist

    eyesCenter = ((eye_1[0] + eye_2[0]) // 2, (eye_1[1] + eye_2[1]) // 2)
    
    # grab the rotation matrix for rotating and scaling the face
    M = cv2.getRotationMatrix2D(eyesCenter, angle, scale)
    # update the translation component of the matrix
    tX = desiredFaceWidth * 0.5
    tY = desiredFaceHeight * desiredLeftEye[1]
    M[0, 2] += (tX - eyesCenter[0])
    M[1, 2] += (tY - eyesCenter[1])

    (w, h) = (desiredFaceWidth, desiredFaceHeight)
    output = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC)

    output = np.array(output)

    output = tf.image.convert_image_dtype(output, dtype=tf.float32)

    output = tf.image.rgb_to_grayscale(output)
    output = tf.tile(output, [1, 1, 3])

    output = tf.clip_by_value(output, 0, 1)

    return output

def predict_remaining_life(img_path):
    detector = mtcnn.MTCNN()
    # Transform the target image and add a batch dimension
    img = load_and_preprocess_image(img_path, detector)
    img = np.expand_dims(img, axis = 0)
    #print(img.shape)
    #plt.imshow(img)
    # Put model into evaluation mode and turn on inference mode
    pred = model.predict(img)
    pred = round(pred[0][0]*100,1)
    
    # Return the prediction dictionary and prediction time 
    return pred


# Create title, description and article strings
title = "Remaining Life Predictor"
description = "A Convolutional Neural Net to predict how many years a person has left to live using their facial image"
article = "Methodology and data explained at [arxiv article](https://arxiv.org/abs/2301.08229)"

# Create the Gradio demo
demo = gr.Interface(fn=predict_remaining_life, # mapping function from input to output
                    inputs=gr.Image(type="filepath"), # what are the inputs?
                    outputs=gr.Number(label="Remaining Life (Year)"),
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch(debug=False, # print errors locally?
            share=False) # generate a publically shareable URL?