import gradio as gr
import torch
import torch.nn.functional as F
from model import create_effnet_b2_model
from timeit import default_timer as timer
from typing import Tuple
import requests
from PIL import Image
from io import BytesIO

# Setup class names
class_names = ['paper', 'rock', 'scissors']

# Model and transforms
model, transforms = create_effnet_b2_model(num_classes=len(class_names))

# Load saved weights
model.load_state_dict(torch.load(f="final_model.pth", map_location=torch.device('cpu')))

# Predict function
def predict(img) -> Tuple[str, float]:

    # Start the timer
    start_time = timer()

    # Transform image and add batch dimension
    transformed_image = transforms(img).unsqueeze(dim=0)

    # Put model into evaluation mode and turn on inference mode
    model.eval()
    with torch.inference_mode():
        # Pass image through model to get raw logits
        target_image_pred = model(transformed_image)
        # Convert prediction logits into prediction probabilities
        target_image_pred_prob = F.softmax(target_image_pred, dim=1)

    # Create prediction label and prediction probability dictionary for each class
    pred_labels_and_probs = {class_names[i]: float(target_image_pred_prob[0][i]) for i in range(len(class_names))}

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

    # Return prediction dictionary and prediction time
    return pred_labels_and_probs, total_time

# Gradio app
title = "Rock paper scissors"
description = "Detects hand sign (rock, paper, scissors) from provided image."
article = "Dataset taken from https://www.kaggle.com/datasets/alexandredj/rock-paper-scissors-dataset?select=scissors. Used transfer learning to get 98% accuracy using EfficientNetB2 model."

# Files from Hugging Face Spaces
files = [
    "examples_1",
    "examples_2",
    "examples_3",
    "examples_4",
    "examples_5",
    "examples_6"
]

# Function to fetch images from Hugging Face Spaces
def fetch_image(filename):
    url = f"https://huggingface.co/spaces/<your-username>/<your-space-name>/resolve/main/{filename}"
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    return img

# Create Gradio examples list
example_list = [[fetch_image(filename)] for filename in files]

# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an image of rock, paper, or scissors"),
    outputs=[gr.Label(num_top_classes=3, label="Predictions"), gr.Number(label="Prediction time")],
    examples=example_list,
    title=title,
    description=description,
    article=article
)

demo.launch(share=True)