import gradio as gr
from transformers import pipeline

#
# 1) Load a broad food classification model
#    This recognizes 100+ food categories (food101).
#
food_classifier = pipeline(
    "image-classification", 
    model="prithivMLmods/Food-101-93M"
)

def estimate_health_and_calories(dish_name: str):

    dish_lower = dish_name.lower()

    # Basic logic: healthy if it sounds like salad, fruit, etc., 
    # less healthy if it sounds like fried, sugary, or dessert.
    
    if any(k in dish_lower for k in ["salad", "fruit", "broccoli", "tomato", "carrot", "spinach"]):
        health = 9
        calories = 80
    elif any(k in dish_lower for k in ["fried", "pizza", "burger", "bacon", "cream", "chips"]):
        health = 3
        calories = 350
    elif any(k in dish_lower for k in ["cake", "pastry", "dessert", "cookie", "chocolate"]):
        health = 2
        calories = 400
    elif any(k in dish_lower for k in ["soup", "stew", "chili"]):
        health = 7
        calories = 150
    elif any(k in dish_lower for k in ["sandwich", "wrap", "taco"]):
        health = 6
        calories = 250
    else:
        # Default fallback
        health = 5
        calories = 200

    return health, calories

def analyze_image(image):

    # Run the model
    outputs = food_classifier(image)

    # Each output item is like {'label': 'omelette', 'score': 0.98}
    # Sort by descending confidence (not strictly necessary if pipeline does so by default)
    outputs = sorted(outputs, key=lambda x: x["score"], reverse=True)
    top_label = outputs[0]["label"]
    top_score = outputs[0]["score"]

    # Confidence threshold to decide if it's "real food" or not
    if top_score < 0.5:
        return "The picture does not depict any food, please upload a different photo."

    # If we pass the threshold, treat the top label as recognized dish
    health_rating, cal_estimate = estimate_health_and_calories(top_label)

    # Build response text
    return (
        f"**Food Identified**: {top_label} (confidence: {top_score:.2f})\n\n"
        f"**Health Rating** (1 = extremely unhealthy, 10 = extremely healthy): **{health_rating}**\n\n"
        f"**Estimated Calories**: ~{cal_estimate} kcal per serving\n\n"
    )


# Build a nice Gradio interface
demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil"),   # PIL image
    outputs="markdown",
    title="Universal Food Recognizer",
    description=(
        "Upload a photo of a dish or ingredient. "
        "We'll attempt to recognize it (from 100+ categories) and rate how healthy it is, "
        "along with a rough calorie estimate. "
        "If no food is detected, you'll see an error message."
    ),
    allow_flagging="never",  # optional: hides the 'flag' button for simpler UI
)

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