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

device = 0 if torch.cuda.is_available() else -1

visual_quest_ans = pipeline("visual-question-answering", model="Salesforce/blip-vqa-base", device=device)
translation_eng_to_ar = pipeline("translation_en_to_ar", model="marefa-nlp/marefa-mt-en-ar", device=device)

def food_recognizer(image):
    result = visual_quest_ans(image=image, question="What is the food or the drink in the image?")
    return result[0]['answer']

def nutrition_info(food):
    data = {"query": food}
    response = requests.post(api_url, headers=headers, json=data)
    return response.json()

def translator(text):
    text = text.strip()
    result = translation_eng_to_ar(text)
    return result[0]['translation_text']

def process_food_result(image, language):
    food_item = food_recognizer(image)
    nutritions_info = nutrition_info(food_item)
    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']
    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')
    
    liquid_keywords = ['juice', 'water', 'milk', 'soda', 'tea', 'coffee']
    is_liquid = any(keyword in food_item.lower() for keyword in liquid_keywords)

    if is_liquid and serving_size != 'Unknown':
        serving_size_text_en = f"{serving_size} mL"
        serving_size_text_ar = f"{serving_size} مل"
    else:
        serving_size_text_en = f"{serving_size} grams"
        serving_size_text_ar = f"{serving_size} جرام"

    if language == "Arabic":
        food_item_ar = translator(food_item)
        return 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>
        """
    else:
        return 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>
        """

def gradio_function(image, language):
    return process_food_result(image, language)

iface = gr.Interface(
    fn=gradio_function,
    inputs=[gr.Image(type="pil", label="Upload an image"),
            gr.Dropdown(choices=["Arabic", "English"], label="Select Language", value="Arabic")],
    outputs=gr.HTML(label="Food and Nutrition Information"),
    # Here, add the custom CSS to style the submit button
    css="""
        .gr-button {
            background-color: #333 !important;
            color: white !important;
            border: none;
        }
    """
)

iface.launch(debug=True)