# -*- coding: utf-8 -*-
"""
Spyder Editor

Use pretrained model to recognize an image.

"""
import os
import numpy as np
import pandas as pd
import PIL
import tensorflow as tf
from tensorflow.keras.applications import EfficientNetV2S
from tensorflow.keras.applications import Xception
from tensorflow.keras.applications import MobileNetV3Small

"""
model = ResNet50(weights='imagenet')

img_path = 'elephant.jpg'
img = keras.utils.load_img(img_path, target_size=(224, 224))
x = keras.utils.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

preds = model.predict(x)
# decode the results into a list of tuples (class, description, probability)
# (one such list for each sample in the batch)
print('Predicted:', decode_predictions(preds, top=3)[0])
"""

# Take a pretrained model, its preprocess input functions, prediction function,
# the image to be predicted.
# Return the decoded prediction.


def pretrained_model(image, model_name, preprocess_fn,
                     decoded_predictions_fn):
    # preprocess with the model's function
    x = preprocess_fn(image)
    # predictions
    pred = model_name.predict(x)
    # decode the results into a list of tuples (class, description, probability)
    decoded_pred = decoded_predictions_fn(pred, top=3)
    # Must put in the [0] first element of the
    # decoded predictions.
    pred_df = pd.DataFrame(
        decoded_pred[0],
        columns=["class", "description", "probability"]
    )
    return pred_df

# Load the specific image size for the model
# path = filepath of the image
# size = size to be loaded


def load_and_resize_image(img_path, size=(224, 224, 3)):
    img = tf.keras.utils.load_img(
        path=img_path, target_size=size)
    x = tf.keras.utils.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    print(f"Shape of x: {x.shape}")

    return x


def predict(img_path):
    # Set the data directory
    data_dir = "images"
    image_list = ["cat.jpg", "mrt_train.jpg",
                  "apples.jpg", "duck.jpg", "daisy.jpg"]

    # Xception:  (299x299)
    # EfficientNetV2S: (384, 384)
    # MobileNetV2: (224, 224, 3)
    # Load the image
    # img_path = os.path.join(data_dir, image_list[4])
    print(f"image path: {img_path}")

    # Efficient Net V2 Small
    img = load_and_resize_image(
        img_path,
        size=(384, 384))
    eff_net = EfficientNetV2S(include_top=True,
                              weights='imagenet',
                              input_shape=(384, 384, 3))
    pred_df_eff = pretrained_model(
        img,
        eff_net,
        tf.keras.applications.efficientnet_v2.preprocess_input,
        tf.keras.applications.efficientnet_v2.decode_predictions,
    )
    print(pred_df_eff[["description", "probability"]])

    # Mobile Net V3 Small
    img = load_and_resize_image(img_path, size=(224, 224))
    mobile_net = MobileNetV3Small(include_top=True,
                                  weights="imagenet",
                                  input_shape=(224, 224, 3))
    pred_df_mob = pretrained_model(
        img,
        eff_net,
        tf.keras.applications.mobilenet_v3.preprocess_input,
        tf.keras.applications.mobilenet_v3.decode_predictions,
    )
    print(pred_df_mob[["description", "probability"]])

    # Xception
    img = load_and_resize_image(img_path, size=(299, 299))
    xcept = Xception(include_top=True,
                     weights="imagenet",
                     input_shape=(299, 299, 3))
    pred_df_xcept = pretrained_model(
        img,
        xcept,
        tf.keras.applications.xception.preprocess_input,
        tf.keras.applications.xception.decode_predictions,
    )
    print(pred_df_xcept[["description", "probability"]])

    # Format the output for gradio label
    dict_eff = dict({
        "Description": pred_df_eff["description"],
        "Probability": pred_df_eff["probability"]
    })
    dict_mob = dict({
        "Description": pred_df_mob["description"],
        "Probability": pred_df_mob["probability"]
    })
    dict_xcept = dict({
        "Description": pred_df_xcept["description"],
        "Probability": pred_df_xcept["probability"]
    })
    return pred_df_eff[["description", "probability"]], \
        pred_df_mob[["description", "probability"]], \
        pred_df_xcept[["description", "probability"]]

 # Main
if __name__ == "__main__":
    predict()