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

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 = [list(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=example_list,
    title=title,
    description=description,
    article=article,
)

demo.lauch()