import streamlit as st
import os
import json
import cv2
import pandas as pd
from ultralytics import YOLO
from PIL import Image

# Configure Streamlit pages
st.set_page_config(
    page_title="NutriDetect: AI-powered Food Detection and Nutritional Insights",
    page_icon="🍴",
      # Set layout to wide for better visualization
)

# Sidebar navigation
st.sidebar.title("Navigation")
page = st.sidebar.radio("Go to", ["Home", "Description", "Try Models"])

# Load nutritional data from the JSON file
def load_nutritional_data(json_path):
    with open(json_path, 'r') as file:
        data = json.load(file)
        for food in data["food"]:
            food["preprocessed_name"] = food["name"].lower().replace(",", "").strip()
        return data

# Load models
@st.cache_resource
def load_models():
    yolov8_model = YOLO("./model/Yolov8m.pt")
    yolov11_model = YOLO("./model/Yolo11m.pt")
    return yolov8_model, yolov11_model

# Global variables for models and nutritional data
models = load_models()
nutritional_json_path = "db.json"
nutritional_data = load_nutritional_data(nutritional_json_path)

# Home Page
if page == "Home":
    st.title("🍴 NutriDetect: Food Detection and Nutritional Insights")
    st.markdown("**Welcome to the NutriDetect App!** Leverage cutting-edge AI to identify food items and provide nutritional details.")
    st.image("back_image.png", use_container_width=True, caption="AI-powered Food Detection")
    st.subheader("🌟 Key Features")
    st.markdown("""
    - **Accurate Food Detection**: Identify food items in images.
    - **Nutritional Insights**: Get calorie, protein, fat, and other details for detected food.
    """)

# Description Page
elif page == "Description":
    st.title("📖 About the NutriDetect App")
    st.markdown("""
    The NutriDetect App uses YOLO-based models for accurate food detection and provides:
    - Bounding boxes for detected food items.
    - Detailed nutritional insights for the detected items.
    """)

# Try Models Page
elif page == "Try Models":
    st.title("🎯 Try the Models")
    uploaded_file = st.file_uploader("Upload an image (jpg, jpeg, png)", type=["jpg", "jpeg", "png"])

    if uploaded_file:
        st.subheader("Step 2: Choose Model")
        model_option = st.selectbox("Select Model", ["YOLOv8m", "YOLO11m"])

        if st.button("🔍 Detect"):
            with st.spinner("Processing..."):
                # Save the uploaded file temporarily
                temp_file_path = "temp_image.jpg"
                with open(temp_file_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())

                # Perform detection
                model = models[0] if model_option == "YOLOv8m" else models[1]
                results = model(temp_file_path)

                # Process detected items
                detected_items = []
                for box in results[0].boxes.data.cpu().numpy():
                    label = int(box[5])  # Get the class label
                    class_name = model.names[label]
                    nutrition_info = next(
                        (food["nutrients"] for food in nutritional_data["food"] if food["preprocessed_name"] == class_name.lower().replace(",", "").strip()),
                        None
                    )
                    if nutrition_info:
                        detected_items.append(
                            {
                                "Food": class_name,
                                "Calories": round(nutrition_info.get("calories", 0), 2),
                                "Protein": round(nutrition_info.get("protein", 0), 2),
                                "Fat": round(nutrition_info.get("fat", 0), 2),
                                "Carbs": round(nutrition_info.get("carbs", 0), 2),
                                "Fiber": round(nutrition_info.get("fiber", 0), 2),
                            }
                        )

                # Handle case with no detected items
                if not detected_items:
                    st.warning("No food items detected, sorry!")
                else:
                    # Convert detected items to DataFrame
                    df = pd.DataFrame(detected_items)

                    # Calculate totals
                    totals = {
                        "Calories": round(df["Calories"].sum(), 2),
                        "Protein": round(df["Protein"].sum(), 2),
                        "Fat": round(df["Fat"].sum(), 2),
                        "Carbs": round(df["Carbs"].sum(), 2),
                        "Fiber": round(df["Fiber"].sum(), 2),
                    }

                    # Display annotated image
                    annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
                    st.image(annotated_image, caption=f"{model_option} Results", use_container_width=True)

                    # Two-column layout for detected items and totals
                    left_col, right_col = st.columns([2, 1])

                    # Left column: Individual food items
                    with left_col:
                        st.header("Detected Food Items")
                        for _, row in df.iterrows():
                            st.subheader(row["Food"])
                            st.write(f"Calories: {row['Calories']} Cal")
                            st.write(f"Protein: {row['Protein']} g")
                            st.write(f"Fat: {row['Fat']} g")
                            st.write(f"Carbs: {row['Carbs']} g")
                            st.write(f"Fiber: {row['Fiber']} g")
                            st.write("---")

                    # Right column: Total nutritional summary
                                        # Right column: Total nutritional summary
                                        # Right column: Total nutritional summary
                    with right_col:
                        st.header("Total Nutritional Summary")
                        st.metric("Total Calories", f"{totals['Calories']} Cal")
                        st.metric("Total Protein", f"{totals['Protein']} g")
                        st.metric("Total Fat", f"{totals['Fat']} g")
                        st.metric("Total Carbs", f"{totals['Carbs']} g")
                        st.metric("Total Fiber", f"{totals['Fiber']} g")
                        
                        # Check for unhealthy items and provide recommendations
                        detected_foods = df["Food"].str.lower().tolist()  # Convert to lowercase for matching
                        unhealthy_items = [item for item in detected_foods if item in ["pizza", "hamburger", "french fries"]]

                        if unhealthy_items:
                            st.warning(
                                f"Detected unhealthy food items: {', '.join(unhealthy_items)}. "
                                "Consider replacing them with fruits and green salad for better health. 🥗🍎"
                            )
                        else:
                            st.success("Balanced meals lead to a healthier life! 😊")