import gradio as gr
import requests
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import torch

# Load the pre-trained Vision Transformer model and image processor
model_name = "google/vit-base-patch16-224"
image_processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# API key for the nutrition information
api_key = '706a5b83a4b567eb535dbc01100bb0f4'
app_id = '8af8f051'

def get_nutrition_info(food_item):
    # Replace with the Nutritionix API endpoint or a suitable API
    nutritionix_url = "https://trackapi.nutritionix.com/v2/natural/nutrients"

    headers = {
        "x-app-id": app_id,
        "x-app-key": api_key,
    }

    response = requests.post(
        nutritionix_url,
        headers=headers,
        json={"query": food_item},
    )

    if response.status_code == 200:
        nutrition_data = response.json()
        return nutrition_data
    else:
        return f"Error: {response.status_code} - {response.text}"

def process_food_image(image):
    # Convert the image to a format suitable for the Vision Transformer
    img = Image.open(image).convert("RGB")
    inputs = image_processor(images=img, return_tensors="pt")

    # Perform prediction
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = torch.argmax(logits, dim=1).item()
    food_item = model.config.id2label.get(predicted_class_idx, "Unknown Food Item")

    # Get nutrition information
    nutrition_info = get_nutrition_info(food_item)

    if isinstance(nutrition_info, dict):
        # Extract and format nutrition data
        nutrition_data = nutrition_info.get("foods", [])[0]  # Avoid duplicate entries
        formatted_data = f"""
        <table>
            <tr><th colspan='2'>Nutrition Facts</th></tr>
            <tr><td>Food Name</td><td>{nutrition_data.get('food_name', 'N/A')}</td></tr>
            <tr><td>Calories</td><td>{nutrition_data.get('nf_calories', 'N/A')} kcal</td></tr>
            <tr><td>Serving Size</td><td>{nutrition_data.get('serving_weight_grams', 'N/A')} g</td></tr>
            <tr><td>Protein</td><td>{nutrition_data.get('nf_protein', 'N/A')} g</td></tr>
            <tr><td>Fat</td><td>{nutrition_data.get('nf_total_fat', 'N/A')} g</td></tr>
            <tr><td>Saturated Fat</td><td>{nutrition_data.get('nf_saturated_fat', 'N/A')} g</td></tr>
            <tr><td>Cholesterol</td><td>{nutrition_data.get('nf_cholesterol', 'N/A')} mg</td></tr>
            <tr><td>Carbohydrates</td><td>{nutrition_data.get('nf_total_carbohydrate', 'N/A')} g</td></tr>
            <tr><td>Sugar</td><td>{nutrition_data.get('nf_sugars', 'N/A')} g</td></tr>
            <tr><td>Fiber</td><td>{nutrition_data.get('nf_dietary_fiber', 'N/A')} g</td></tr>
            <tr><td>Potassium</td><td>{nutrition_data.get('nf_potassium', 'N/A')} mg</td></tr>
        </table>
        """
        return formatted_data
    else:
        return nutrition_info

# Gradio Interface
iface = gr.Interface(
    fn=process_food_image,
    inputs=gr.Image(type="filepath"),
    outputs="html",
    title="Nutrition Agent App",
    description="Upload a picture of food, and get accurate nutrition information in a formatted table."
)

if __name__ == "__main__":
    iface.launch()