### 1. Imports and class names setup ###
import gradio as gr
import os
import torch

from model import create_effnetb2_model
from timeit import default_timer as timer
from typing import Tuple, Dict

#calorie count
import requests

def get_calorie_count(food_name):
    # Use USDA FoodData Central Search API
    api_key = "8Kw59j70SiLY9BWuBBOYOQkLzwcsXVFEJO4pVvpf"
    search_url = f"https://api.nal.usda.gov/fdc/v1/foods/search"
    params = {
        "api_key": api_key,
        "query": food_name,
        "pageSize": 1,
    }
    try:
        response = requests.get(search_url, params=params)
        response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
        data = response.json()
        if "foods" in data and data["foods"]:
            food = data["foods"][0]
            for nutrient in food.get("foodNutrients", []):
                if nutrient.get("nutrientName", "").lower() == "energy" and nutrient.get("unitName", "").lower() == "kcal":
                    return f'{nutrient.get("value")} kcal per 100g'
        return "Calorie info not found"
    except requests.exceptions.RequestException as e:
        return f"API Error: {e}" # More specific error handling
    except Exception as e:
        return f"General Error: {e}"

# Pre-fetch calorie values for all food classes
class_calories = {}

def preload_calories():
    with open('class_names.txt', 'r') as f:
        for food_name in f:
            food_name = food_name.strip() # Remove leading/trailing whitespace
            if food_name: # Ensure the line is not empty
                # Replace underscores with spaces
                formatted_food_name = food_name.replace('_', ' ')
                #print(f"Searching for: {formatted_food_name}") # Added for debugging
                class_calories[food_name] = get_calorie_count(formatted_food_name)

preload_calories()

def get_calorie_count(food_name):
    return class_calories.get(food_name, "Calorie info not found")

# Setup class names
with open("class_names.txt", "r") as f: # reading them in from class_names.txt
    class_names = [food_name.strip() for food_name in  f.readlines()]

### 2. Model and transforms preparation ###

# Create model
effnetb2, effnetb2_transforms = create_effnetb2_model(
    num_classes=101, # could also use len(class_names)
)

# Load saved weights
effnetb2.load_state_dict(
    torch.load(
        f="pretrained_effnetb2_feature_extractor_food101_20_percent.pth",
        map_location=torch.device("cpu"),  # load to CPU
    )
)

### 3. Predict function ###

# Create predict function
def predict(img) -> Tuple[Dict, float]:
    """Transforms and performs a prediction on img and returns prediction and time taken.
    """
    # Start the timer
    start_time = timer()

    # Transform the target image and add a batch dimension
    img = effnetb2_transforms(img).unsqueeze(0)

    # Put model into evaluation mode and turn on inference mode
    effnetb2.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        pred_probs = torch.softmax(effnetb2(img), dim=1)

    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

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

    # Return the prediction dictionary and prediction time
    # return pred_labels_and_probs, pred_time
    top_food = max(pred_labels_and_probs, key=pred_labels_and_probs.get)
    calorie_count = get_calorie_count(top_food)
    return pred_labels_and_probs, pred_time, calorie_count

### 4. Gradio app ###

# Create title, description and article strings
title = "Food Vision 🍔👁"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of food into 101 different classes"
article = "Created at : https://github.com/rageya/food-vision"

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

# Create Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=5, label="Predictions"),
        gr.Number(label="Prediction time (s)"),
        gr.Textbox(label="Estimated Calories (per 100g)"),
    ],
    examples=example_list,
    title=title,
    description=description,
    article=article,
    #api_name="predict"
)

# Launch the app!
# Launch the app with API explicitly enabled!
demo.launch(share=True)