import os
import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer
from sentence_transformers import SentenceTransformer, util
import torch

# Prevent parallelism warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Load the models and tokenizer using Streamlit's updated caching method
@st.cache_resource
def load_models():
    generation_model = T5ForConditionalGeneration.from_pretrained('t5-small')
    generation_tokenizer = T5Tokenizer.from_pretrained('t5-small')
    sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
    return generation_model, generation_tokenizer, sentence_model

generation_model, generation_tokenizer, sentence_model = load_models()

# Predefined meals with calorie information
# Updated predefined meals with higher calorie options
meals = {
    "Breakfast": [
        {"meal": "Oatmeal with fruits", "calories": 250},
        {"meal": "Scrambled eggs with toast", "calories": 350},
        {"meal": "Protein Smoothie", "calories": 300},
        {"meal": "Avocado toast with poached egg", "calories": 320},
        {"meal": "Greek yogurt with granola and berries", "calories": 280},
        {"meal": "Pancakes with honey", "calories": 400},
        {"meal": "Veggie omelette", "calories": 300},
        {"meal": "Cereal with milk", "calories": 400}
    ],
    "Lunch": [
        {"meal": "Grilled chicken salad", "calories": 400},
        {"meal": "Vegetable stir-fry with tofu", "calories": 350},
        {"meal": "Chicken wrap", "calories": 500},
        {"meal": "Turkey sandwich on whole grain bread", "calories": 450},
        {"meal": "Quinoa salad with mixed greens", "calories": 375},
        {"meal": "Fish tacos", "calories": 480},
        {"meal": "Lentil soup with a side salad", "calories": 420},  
        {"meal": "Pasta Alfredo with garlic bread", "calories": 1100},  # High-calorie option
        {"meal": "Buffalo wings with ranch dressing", "calories": 950}  # High-calorie option
    ],
    "Dinner": [
        {"meal": "Salmon with steamed vegetables", "calories": 600},
        {"meal": "Pasta with tomato sauce", "calories": 500},
        {"meal": "Vegetarian chili", "calories": 450},
        {"meal": "Beef stir-fry with bell peppers", "calories": 550},
        {"meal": "Chicken curry with brown rice", "calories": 650},
        {"meal": "Grilled shrimp with quinoa", "calories": 580},
        {"meal": "Stuffed bell peppers with ground turkey", "calories": 525},
        {"meal": "Ribeye steak with mashed potatoes", "calories": 1300},  # High-calorie option
        {"meal": "BBQ pork ribs with cornbread", "calories": 1400},  # High-calorie option
        {"meal": "Seafood platter with garlic butter", "calories": 1250}  # High-calorie option
    ],
    "Snacks": [
        {"meal": "Apple with peanut butter", "calories": 200},
        {"meal": "Greek yogurt with honey", "calories": 150},
        {"meal": "Granola bar", "calories": 250},
        {"meal": "Handful of almonds", "calories": 180},
        {"meal": "Hummus with carrot sticks", "calories": 100},
        {"meal": "Cheese and crackers", "calories": 200},
        {"meal": "Protein shake", "calories": 220}
    ]
}


# Function to suggest meals based on the target calorie goal
def suggest_meals(target_calories):
    suggestions = {}
    total_calories = 0
    remaining_calories = target_calories

    # Prioritize essential meals
    essential_meal_categories = ["Breakfast", "Lunch", "Dinner"]
    snack_category = "Snacks"

    for category in essential_meal_categories:
        # Sort meals by calories in descending order to prioritize higher calorie meals
        available_meals = sorted(meals[category], key=lambda x: -x["calories"])
        selected_meal = None

        for meal in available_meals:
            if meal["calories"] <= remaining_calories:
                selected_meal = meal
                break

        if selected_meal:
            suggestions[category] = f"{selected_meal['meal']} containing {selected_meal['calories']} calories"
            total_calories += selected_meal["calories"]
            remaining_calories -= selected_meal["calories"]

    # Add multiple snacks to use up the remaining calories efficiently
    available_snacks = sorted(meals[snack_category], key=lambda x: -x["calories"])
    snack_descriptions = []

    for snack in available_snacks:
        if snack["calories"] <= remaining_calories:
            snack_descriptions.append(f"{snack['meal']} containing {snack['calories']} calories")
            total_calories += snack["calories"]
            remaining_calories -= snack["calories"]

            if remaining_calories < 100:  # Stop adding snacks if the remaining calories are too low
                break

    return suggestions, snack_descriptions, total_calories

# Streamlit interface
st.title("Calorie-Based Meal Plan Bot")
st.write("Enter your daily calorie goal, and I will suggest a balanced meal plan for the day. Minimum calorie goal is 800 kcal for a sustainable diet.")

# Ensure the user inputs a calorie target without a default value
target_calories = st.number_input("Enter your daily calorie target (in kcal):", min_value=800, max_value=4500, value=None)

if target_calories:
    suggested_meals, snack_descriptions, total_calories = suggest_meals(target_calories)
    
    st.subheader("Suggested Meals:")
    for category, meal_description in suggested_meals.items():
        st.write(f"{category}: {meal_description}")
    
    
    if(len(snack_descriptions) == 0):
        st.write("Snacks: Sorry. No snacks for you 🙁")
    else:
        st.write("Snacks:")
        for snack in snack_descriptions:
            st.write(f"- {snack}")
    
    st.write(f"Total Calories: {total_calories} calories")
    
    if total_calories < target_calories:
        st.write("You can add more snacks or meals to meet your target!")
    elif total_calories > target_calories:
        st.write("Consider reducing meal portions to stay within your target!")