# Import necessary libraries
import gradio as gr
import wget
from transformers import pipeline
import requests
import torch

# Nutritionix API setup
api_url = "https://trackapi.nutritionix.com/v2/natural/nutrients"

# App ID, App Key provided by Nutritionix
headers = {
    "x-app-id": "dd773727",
    "x-app-key": "86f278fc4c7f276c386f280848acf3e6",
}

# Load the Models

# Check if a GPU is available
device = 0 if torch.cuda.is_available() else -1

# Load the BLIP VQA Model (Recognize the food)
visual_quest_ans = pipeline("visual-question-answering", model="Salesforce/blip-vqa-base", device=device)

# Load the Translation Model (English to Arabic)
translation_eng_to_ar = pipeline("translation_en_to_ar", model="marefa-nlp/marefa-mt-en-ar", device=device)

# Function to recognize food from the image using the VQA model
def food_recognizer(image):
    # Pass the image and the question to the model to identify the food on the image
    result = visual_quest_ans(image=image, question="What is the food or the drink in the image?")
    return result[0]['answer']

# Function to fetch nutritional information from Nutritionix API
def nutrition_info(food):
    # Prepare the data for the API request
    data = {
        "query": food
    }

    # Send a POST request to the Nutritionix API with the food item
    response = requests.post(api_url, headers=headers, json=data)

    # Get the nutritional information in JSON format
    nutritions = response.json()
    return nutritions

# Function to translate text from English to Arabic with preprocessing
def translator(text):
    text = text.strip()  # Remove leading/trailing spaces
    result = translation_eng_to_ar(text) # Use the translation model to translate the text
    result = result[0]['translation_text']
    return result

# Function to process food recognition and get nutrition info
def process_food_result(image, language):
    # Recognize the food item in the uploaded image
    food_item = food_recognizer(image)
    
    # Fetch nutritional information for the recognized food item
    nutritions_info = nutrition_info(food_item)

    # Extract nutritional information
    food_info = nutritions_info['foods'][0]
    calories = food_info['nf_calories']
    protein = food_info['nf_protein']
    carbs = food_info['nf_total_carbohydrate']
    fat = food_info['nf_total_fat']
    # Use 'Unknown' if value is not available
    sugars = food_info.get('nf_sugars', 'Unknown')
    fiber = food_info.get('nf_dietary_fiber', 'Unknown')
    sodium = food_info.get('nf_sodium', 'Unknown')
    serving_size = food_info.get('serving_weight_grams', 'Unknown')
    
    # Identify if the food item is a liquid (simple check for common drink categories)
    liquid_keywords = ['juice', 'water', 'milk', 'soda', 'tea', 'coffee']
    is_liquid = any(keyword in food_item.lower() for keyword in liquid_keywords)

    # Convert serving size to milliliters if it's a liquid
    if is_liquid and serving_size != 'Unknown':
        serving_size_ml = serving_size  # Assume 1 gram ≈ 1 milliliter for liquids
        serving_size_text_en = f"{serving_size_ml} mL"
        serving_size_text_ar = f"{serving_size_ml} مل"
    else:
        serving_size_text_en = f"{serving_size} grams"
        serving_size_text_ar = f"{serving_size} جرام"

    # Generate output in the selected language
    if language == "Arabic":
        # Translate the food item name to Arabic
        food_item_ar = translator(food_item)
        output_ar = f"""
        <div style='direction: rtl; text-align: right;'>
            <b>الطعام</b>: {food_item_ar}<br>
            <b>حجم الحصة</b>: {serving_size_text_ar}<br>
            <b>السعرات الحرارية</b>: {calories} كيلو كالوري<br>
            <b>البروتين</b>: {protein} جرام<br>
            <b>الكربوهيدرات</b>: {carbs} جرام<br>
            <b>السكر</b>: {sugars} جرام<br>
            <b>الألياف</b>: {fiber} جرام<br>
            <b>الصوديوم</b>: {sodium} مجم<br>
            <b>الدهون</b>: {fat} جرام
        </div>
        """
        return output_ar
    else:
       # For English output
        output_en = f"""
        <div style='text-align: left;'>
            <b>Food</b>: {food_item}<br>
            <b>Serving Size</b>: {serving_size_text_en}<br>
            <b>Calories</b>: {calories} kcal<br>
            <b>Protein</b>: {protein}g<br>
            <b>Carbohydrates</b>: {carbs}g<br>
            <b>Sugars</b>: {sugars}g<br>
            <b>Fiber</b>: {fiber}g<br>
            <b>Sodium</b>: {sodium}mg<br>
            <b>Fat</b>: {fat}g
        </div>
        """
        return output_en


# Gradio interface function
def gradio_function(image, language):
    # Call the process_food_result function to get the output
    result = process_food_result(image, language)
    return result

# Define URLs of example images
image_urls = [
    "https://raw.githubusercontent.com/Abdulrahman078/ML_Datasets-Imgs-Vids/main/close-up-delicious-pizza.jpg",
    "https://raw.githubusercontent.com/Abdulrahman078/ML_Datasets-Imgs-Vids/main/assorted-desserts-with-chocolate-frosted-pink-glazed-sprinkles.jpg",
    "https://raw.githubusercontent.com/Abdulrahman078/ML_Datasets-Imgs-Vids/main/fried-fish-with-cranberries-wooden-board.jpg",
    "https://raw.githubusercontent.com/Abdulrahman078/ML_Datasets-Imgs-Vids/main/glass-water.jpg"
]

# Download the images and use their paths
example_images = [wget.download(url) for url in image_urls]
examples = [[img] for img in example_images]


# Setup the Gradio interface
iface = gr.Interface(
    fn=gradio_function, # Function to call
    inputs=[gr.Image(type="pil", label="Upload an image"), # Input: Image (in PIL format)
            gr.Dropdown(choices=["Arabic", "English"], label="Select Language", value="Arabic")], # Input: Dropdown for language selection
    outputs=gr.HTML(label="Food and Nutrition Information"), # Output: HTML for displaying nutrition info
    title="Bilingual Food Recognition and Nutrition Info Tool", # Title of the Gradio interface
    description="Upload an image of food, and the tool will recognize it and provide nutritional information in both English or Arabic languages.", # Description of the tool
    examples=examples  # Add examples with the image and language
)

# Launch the Gradio interface with debug mode enabled
iface.launch(debug=True)