### 1. Imports and class names setup ###
from model import create_model
import pandas as pd
import torch
from typing import Tuple, Dict
from timeit import default_timer as timer
import gradio as gr
import os
import numpy as np

### 1. Create a Dictionary for Dog Breeds
labels_csv = pd.read_csv('./labels.csv')
labels = labels_csv['breed']
labels = np.array(labels)
unique_labels = np.unique(labels)

unique_labels = [' '.join([word.capitalize() for word in label.split('_')]) for label in unique_labels]

### 2. Model and transforms preparation ###
model, model_transforms = create_model(num_classes=len(unique_labels))
model = torch.compile(model)


# Load save weights
model.load_state_dict(torch.load(f='./convnext_model.pth', map_location='cpu',weights_only=True))

# 3. Predict Function

def predict(img) -> Tuple[Dict[str, float], str]:
    """
    Predicts the class probabilities for a given image using a pre-trained model.

    Args:
        img: A PIL image to be predicted.

    Returns:
        A tuple containing:
        - A formatted string displaying class labels and their respective probabilities.
        - The time taken for inference in seconds as a string.
    """
    # Start a timer
    start_time = timer()

    # Put the model into evaluation mode and disable gradient computation
    model.eval()
    with torch.inference_mode():
        
        # Transform the input image for use with the model
        img = model_transforms(img).unsqueeze(dim=0)
        
        # Pass transformed image through the model
        pred_logit = model(img)

    # Turn prediction logits into probabilities
    pred_prob = torch.softmax(pred_logit, dim=1)
    
    pred_label = torch.argmax(pred_prob, dim=1)

    # Map probabilities to class labels
    prediction = unique_labels[pred_label]
    probabilities = {unique_labels[i]: pred_prob[0, i].item() for i in range(len(unique_labels))}

    # Calculate the time taken
    end_time = timer()
    inference_time = end_time - start_time

    # Return predictions as a dictionary and inference time
    return probabilities, f"{inference_time:.4f} seconds"
### 4. Gradio app ###


# Create title, description and article
title = "Dogvision 🐶"
description = "A [ConvNeXt Tiny](https://pytorch.org/vision/stable/models/generated/torchvision.models.convnext_tiny.html#torchvision.models.convnext_tiny) Computer Vision Model To Classify 120 Dog Breeds 🐩 Ranging fro A Labrador 🐕 to A German Shepherd! 🐕‍🦺"
article = "Created with 🤎 (and a mixture of mathematics, statistics, and tons of calculations 👩🏽‍🔬) by Arpit Vaghela [GitHub](https://github.com/magnifiques)"

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

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type='pil'),
                    outputs=[
                    gr.Label(num_top_classes=3, label="Top Predictions"),  # Display top predictions with probabilities
                    gr.Textbox(label="Prediction Time (s)")  # Display inference time
                    ],
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)

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