import os
from timeit import default_timer as timer
from typing import Tuple
from pathlib import Path
from PIL import Image

import gradio as gr
import torch
from torch import nn
from torchvision import transforms

from model import create_effnetb2_model

class_names = ["pizza", "steak", "sushi"]
device = "cpu"

# Create model
effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=len(class_names))

# Load saved weights
effnetb2.load_state_dict(torch.load("effnetb2.pth", map_location=torch.device(device)))

# Define predict function
def predict(img: Image) -> Tuple[dict, float]:
    """Uses EffnetB2 model to transform and predict on img. Returns prediction
    probabilities and time taken.
    
    Args:
      img (PIL.Image): Image to predict on.
    
    Returns:
      A tuple (pred_labels_and_probs, pred_time), where pred_labels_and_probs
      is a dict mapping each class name to the probability the model assigns to
      it, and pred_time is the time taken to predict (in seconds).
    """
    start_time = timer()
    img = effnetb2_transforms(img).unsqueeze(0)
    effnetb2.eval()
    with torch.inference_mode():
        pred_probs = torch.softmax(effnetb2(img), dim=1)
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i])
                             for i in range(len(class_names))}
    pred_time = round(timer() - start_time, 4)
    return pred_labels_and_probs, pred_time

# Initialize Gradio app
title = "FoodVision Mini"
description = "EfficientNetB2 feature extractor to classify images of food as pizza, steak, or sushi."
article = "From the [Zero to Mastery PyTorch tutorial](https://www.learnpytorch.io/09_pytorch_model_deployment/)"
examples = [[example] for example in Path("examples").glob("*.jpg")]

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(num_top_classes=3, label="Predictions"),
             gr.Number(label="Prediction time (s)")], 
    examples=examples,
    title=title,
    description=description,
    article=article,
)

demo.launch()