from keras.models import load_model
from PIL import Image
import numpy as np
import requests
from io import BytesIO
import timeit
import gdown
import json
import os

# Define the class labels for the model's output

# output_class = ["Plastic bottle/Can to deposit in Supermarkets", "Big Cardboard bin", "Unrecyclable garbage", "Glass - Purple bin", "Organic waste - Composter", "Grocery Packages - Orange bin", "Paper - Blue bin"]

with open('Bins.json', 'r') as f:
    bins_dict = json.load(f)
    output_class = bins_dict["output_class"]

url = "https://drive.google.com/uc?id=1TJSpf9wJYZSqIm5ytqXkU8RLuKJmmHHa"
output = 'vgg16_model_after_one_epoch.h5'

# check if path to file exist
if not os.path.exists(output):
    gdown.download(url, output, quiet=False)

# Load the pre-trained model
model = load_model(output)


def predict(img_input, is_url=False):
    """
    Function to predict the class labels and probabilities of an input image.

    Args:
        file (str): Path to the input image file or URL.
        is_url (bool): Indicates whether the file argument is a URL or local file path.

    Returns:
        dict: A dictionary containing the predicted class labels as keys and
              their corresponding probabilities as values.
    """
    if isinstance(img_input, str):
        if is_url is True:
            # Download the image from the URL
            response = requests.get(img_input)
            img = Image.open(BytesIO(response.content))
        else:
            # Open the image file from the local file path
            img = Image.open(img_input)
    else:
        img = img_input

    # Resize the image to match the input size expected by the model (224x224)
    resized_image = img.resize((224, 224))

    # Normalize the pixel values to be between 0 and 1
    normalized_image = np.array(resized_image) / 255.0

    # Expand dimensions to match the input shape expected by the model
    pred_image = np.expand_dims(normalized_image, axis=0)

    # Make predictions using the model
    predicted_array = model.predict(pred_image)

    # Create a dictionary of predicted class labels and probabilities
    pred_labels_and_probs = {
        output_class[i]: float(predicted_array[0][i])
        for i in range(len(output_class))
    }

    return pred_labels_and_probs


def image_identity(image):
    """
    Take an image as input and return prediction information and an image of the predicted waste bin.

    Parameters:
    image (np.array): An image in numpy array format.

    Returns:
    bin_image1 (np.array): An image of the predicted waste bin.
    predicted_class1 (str): The class of the predicted bin.
    predicted_prob1 (float): The probability of the predicted class.
    predicted_class2 (str): The second most likely class.
    predicted_prob2 (float): The probability of the second most likely class.
    prediction_time (str): The time taken to make the prediction.
    """
    image = Image.fromarray((image * 255).astype(np.uint8))

    # Load the classification and image data from the JSON file
    with open('Bins.json', 'r') as f:
        bins_dict = json.load(f)

    # Retrieve the class labels and bin images from the dictionary
    output_class = bins_dict["output_class"]
    bin_images = bins_dict["bin_images"]

    # Record the start time before making predictions
    start_time = timeit.default_timer()

    # Make predictions on the provided image
    predicted_class1, predicted_prob1, predicted_class2, predicted_prob2 = check_item(image)

    # Record the end time after making predictions and calculate the total time taken
    end_time = timeit.default_timer()
    pred_time = round(end_time - start_time, 2)
    prediction_time = f"{pred_time} seconds"

    # Retrieve the image URL of the predicted bin
    bin_1 = output_class.index(predicted_class1)
    bin_image_url1 = bin_images[bin_1]
    response1 = requests.get(bin_image_url1)

    # Open the image and convert to a numpy array for returning
    bin_image1 = Image.open(BytesIO(response1.content))
    bin_image1 = np.array(bin_image1)

    return bin_image1, predicted_class1, predicted_prob1, predicted_class2, predicted_prob2, prediction_time


def check_item(file_path, is_url=False):
    """
    Function to make predictions on an input image and compare the results with the expected classes.

    Args:
        file_path (str): Path to the input image file or URL.
        is_url (bool): Indicates whether the file argument is a URL or local file path.
    """
    # Obtain the predicted class labels and probabilities
    prediction_dict = predict(file_path, is_url)

    # Get the index of the highest prediction
    ind = np.argmax(list(prediction_dict.values()))
    predicted_class = list(prediction_dict.keys())[ind]
    predicted_prob = list(prediction_dict.values())[ind]
    predicted_prob = f"{predicted_prob:.2f}"
    print(f"1st Prediction probabilities: {predicted_class} {predicted_prob}")

    # Get the index of the second highest prediction
    sorted_probs = np.sort(list(prediction_dict.values()))
    ind2 = np.argsort(list(prediction_dict.values()))[-2]
    predicted_class2 = list(prediction_dict.keys())[ind2]
    predicted_prob2 = sorted_probs[-2]
    predicted_prob2 = f"{predicted_prob2:.2f}"
    print(f"2nd Prediction probabilities: {predicted_class2} {predicted_prob2}")
    return predicted_class, predicted_prob, predicted_class2, predicted_prob2


if __name__ == "__main__":
    # Prompt the user to enter the file path to the input image
    file_path = input("Enter the file path or URL to the input image: ")

    # Prompt the user to indicate if the input is a URL or a local file path
    is_url = input("Is the input a URL? (y/n): ").lower() == "y"

    # Measure the time taken to make predictions
    start_time = timeit.default_timer()

    # Make predictions and compare with expected results
    check_item(file_path, is_url)

    # Calculate and print the prediction time
    end_time = timeit.default_timer()
    pred_time = round(end_time - start_time, 2)
    print("Prediction time:", pred_time, "seconds")