import gradio as gr
import os
import tensorflow

from typing import Tuple, Dict
import tensorflow as tf
import numpy as np
from PIL import Image
from timeit import default_timer as timer


# Setup class names
with open('class_names.txt') as f:
    class_names = [breed.strip() for breed in f.readlines()]

# Load the pre-trained EfficientNetV2 model
effnet = tensorflow.keras.models.load_model('demo/dog_breed_classifier/dog_breed_effnet_augmentation.h5')

# Get the preprocessing function for EfficientNetV2
effnet_preprocess_input = tensorflow.keras.applications.efficientnet_v2.preprocess_input

# Create examples list from "examples/" directory
example_list = [["examples/" + example] for example in os.listdir("examples")]

## `predict` function
def predict(img) -> Tuple[Dict, float]:
    """
    Transforms and performs a prediction on an image and returns prediction and time taken.

    Args:
        image_path (str): Path to the input image.

    Returns:
        Tuple[Dict, float]: A tuple containing a dictionary of class labels and prediction probabilities
        and the prediction time.
    """
    # Start the timer
    start_time = timer()

    # Open the image using PIL
    img = img.resize((224, 224))  # Resize the image to the model's expected input size

    # Convert the image to a NumPy array
    x = np.array(img)
    x = effnet_preprocess_input(x)

    # Add a batch dimension
    x = tf.expand_dims(x, axis=0)

    # Pass the image through the model
    predictions = effnet(x)

    top_classes_indices = np.argsort(predictions[0])[::-1][:3]  # Get the indices of top 3 classes
    top_classes = [class_names[i] for i in top_classes_indices]  # Get the class names of top 3 classes
    top_probabilities = [predictions[0][index] for index in top_classes_indices] * 100  # Get the probabilities of top 3 classes

    # Create a dictionary of class labels and prediction probabilities
    pred_labels_and_probs = {top_classes[i]: float(top_probabilities[i]) for i in range(len(top_classes_indices))}

    # Calculate the prediction time
    pred_time = round(timer() - start_time, 5)

    # Return the prediction dictionary and prediction time
    return pred_labels_and_probs, pred_time


# Create title, description, and article strings
title = "🐶 Dog Breeds Classifier 🐾"
description = "🚀 An EfficientNetV2S feature extractor computer vision model to classify images of 120 different breeds. 📸"
article = "🌟 Created at [GitHub](https://github.com/adinmg/dog_breed_classifier)."

# Create the Gradio demo
demo = gr.Interface(
    fn=predict,  # mapping function from input to output
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=3, label="Predictions"),  # what are the outputs?
        gr.Number(label="Prediction time (s)"),
    ],
    examples=example_list,
    title=title,
    description=description,
    article=article,
)

# Launch the demo!
demo.launch()