import streamlit as st
from dotenv import load_dotenv
from langchain_groq import ChatGroq
import os
import importlib.util

# Set page configuration
st.set_page_config(
    page_title="Wellness Planner",
    page_icon="🌿",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main {
        padding: 2rem;
    }
    .stApp {
        background-color: #f8f9fa;
    }
    .css-18e3th9 {
        padding-top: 2rem;
    }
    .stTabs [data-baseweb="tab-list"] {
        gap: 2px;
    }
    .stTabs [data-baseweb="tab"] {
        height: 50px;
        white-space: pre-wrap;
        background-color: #f0f2f6;
        border-radius: 4px 4px 0 0;
        gap: 1px;
        padding-top: 10px;
        padding-bottom: 10px;
    }
    .stTabs [aria-selected="true"] {
        background-color: #4CAF50 !important;
        color: white !important;
    }
    div.stButton > button:first-child {
        background-color: #4CAF50;
        color: white;
        font-weight: bold;
        padding: 0.5rem 1rem;
        border-radius: 0.5rem;
        border: none;
    }
    div.stButton > button:hover {
        background-color: #45a049;
        border: none;
    }
    .result-container {
        background-color: white;
        padding: 2rem;
        border-radius: 0.5rem;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .info-box {
        background-color: #e8f4f8;
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
    h1, h2, h3 {
        color: #2E7D32;
    }
</style>
""", unsafe_allow_html=True)

# Check if team_info module exists and import it
try:
    import team_info
    has_team_info = True
except ImportError:
    has_team_info = False

# Load environment variables
load_dotenv()

# Initialize the model
model = ChatGroq(model="llama3-8b-8192")

# Function to calculate daily calorie requirements
def calculate_calorie_requirements(age, gender, weight, height, fitness_goal):
    if gender == "Male":
        bmr = 10 * weight + 6.25 * height - 5 * age + 5
    else:
        bmr = 10 * weight + 6.25 * height - 5 * age - 161

    if fitness_goal == "Weight Loss":
        return bmr * 1.2
    elif fitness_goal == "Weight Gain":
        return bmr * 1.5
    else:
        return bmr * 1.375

# Function to generate the plan
def generate_plan_with_prompt(metrics, prompt_template):
    prompt = prompt_template.format(**metrics)
    response = model.invoke(prompt)
    return response

# Function to format the response neatly
def format_plan(response):
    try:
        content = response.content
        sections = content.split("\n\n")
        formatted = ""
        for section in sections:
            formatted += f"{section.strip()}\n\n"
        return formatted
    except Exception as e:
        return f"Error formatting plan: {e}"

# Ayurvedic prompt template
ayurveda_prompt_template = """
You are a health expert specialized in both modern medicine and Ayurveda. Generate a personalized weekly diet and exercise plan for {name}, a {age}-year-old {gender} with a BMI of {bmi} ({health_status}).

Fitness Goal: {fitness_goal}.
Daily Calorie Requirement: {daily_calories} kcal.
Dietary Preference: {dietary_preference}.
Food Allergies: {food_allergies}.
Local Cuisine: {local_cuisine}.
Month: {month}.
Ayurvedic Consideration: True

Plan should include:
1. A daily diet plan with meal timings, calorie details, and meal alternatives.
2. Exercise routines based on goals, incorporating cardio, strength, and flexibility.
3. Dynamic plan adjustments based on month and local cuisine preferences.
4. Wearable integration for tracking steps, heart rate, and calorie burn.
5. Progress monitoring for daily calorie burn and weight tracking.
6. **Food Delivery Integration**:
   - Meal suggestions based on diet plans.
   - Integration with food delivery platforms (Uber Eats, DoorDash).
   - Searching menu items that fit calorie and dietary preferences.
   - Multi-restaurant meal aggregation for complete diet fulfillment.
   - Location-based meal recommendations.
   - Customizable meal delivery schedules.
7. **Ayurvedic Elements**:
   - Dosha assessment based on provided metrics.
   - Ayurvedic dietary recommendations aligned with your dosha type.
   - Seasonal herbs and spices for balance.
   - Daily routines (Dinacharya) for optimal health.
   - Natural remedies complementing modern medicine.

Provide a detailed plan for each weekday: Monday, Tuesday, Wednesday, Thursday, Friday, Saturday, Sunday.

Return output as:
Day: {{weekday}}
  - Breakfast: Time, Description, Calories, Ayurvedic Properties
  - Lunch: Time, Description, Calories, Ayurvedic Properties
  - Snacks: Time, Description, Calories, Ayurvedic Properties
  - Dinner: Time, Description, Calories, Ayurvedic Properties
  - Exercise: Description, Duration, Dosha Benefits
  - Wearable Tracking: Steps, Heart Rate, Calories Burned
  - Progress Monitoring: Daily calorie intake vs. burn.
  - Food Delivery: Suggested meal items and delivery options.
  - Ayurvedic Tips: Daily wellness practices based on dosha.
"""

# Regular prompt template
regular_prompt_template = """
You are a health expert. Generate a personalized weekly diet and exercise plan for {name}, a {age}-year-old {gender} with a BMI of {bmi} ({health_status}).

Fitness Goal: {fitness_goal}.
Daily Calorie Requirement: {daily_calories} kcal.
Dietary Preference: {dietary_preference}.
Food Allergies: {food_allergies}.
Local Cuisine: {local_cuisine}.
Month: {month}.

Plan should include:
1. A daily diet plan with meal timings, calorie details, and meal alternatives.
2. Exercise routines based on goals, incorporating cardio, strength, and flexibility.
3. Dynamic plan adjustments based on month and local cuisine preferences.
4. Wearable integration for tracking steps, heart rate, and calorie burn.
5. Progress monitoring for daily calorie burn and weight tracking.
6. **Food Delivery Integration**:
   - Meal suggestions based on diet plans.
   - Integration with food delivery platforms (Uber Eats, DoorDash).
   - Searching menu items that fit calorie and dietary preferences.
   - Multi-restaurant meal aggregation for complete diet fulfillment.
   - Location-based meal recommendations.
   - Customizable meal delivery schedules.

Provide a detailed plan for each weekday: Monday, Tuesday, Wednesday, Thursday, Friday, Saturday, Sunday.

Return output as:
Day: {{weekday}}
  - Breakfast: Time, Description, Calories
  - Lunch: Time, Description, Calories
  - Snacks: Time, Description, Calories
  - Dinner: Time, Description, Calories
  - Exercise: Description, Duration
  - Wearable Tracking: Steps, Heart Rate, Calories Burned
  - Progress Monitoring: Daily calorie intake vs. burn.
  - Food Delivery: Suggested meal items and delivery options.
"""

# Main App Layout
def main():
    # App Header
    col1, col2 = st.columns([1, 3])
    with col1:
        st.image("https://img.icons8.com/color/96/000000/yoga.png", width=80)
    with col2:
        st.title("AI-Based Wellness Planner")
        st.subheader("Personalized Diet & Exercise Plan with Ayurvedic Insights")
    
    st.markdown("---")
    
    # Sidebar content
    with st.sidebar:
        st.image("https://img.icons8.com/color/96/000000/natural-food.png", width=50)
        st.header("About this Project")
        st.markdown("""
        <div class="info-box">
        This app creates personalized wellness plans by integrating Ayurveda with modern medical science for comprehensive health insights.
        </div>
        """, unsafe_allow_html=True)
        
        if has_team_info:
            st.subheader("👥 Team Members")
            team_members = [
                ("Aditi Soni", "https://www.linkedin.com/in/aditi-soni-259813285/"),
                ("Bhumika Patel", "https://www.linkedin.com/in/bhumika-patel-ml/"),
                ("Aditi Lakhera", "https://www.linkedin.com/in/aditi-lakhera-b628802bb/"),
                ("Anushri Tiwari", "https://www.linkedin.com/in/anushri-tiwari-916494300")
            ]
            
            for name, link in team_members:
                st.markdown(f"- [{name}]({link})")
            
            st.subheader("🔗 Project Links")
            st.markdown("[GitHub Repository](https://github.com/Abs6187/AI_Health_v2)")
            st.markdown("[Presentation](https://github.com/Abs6187/AI_Health_v2/blob/main/HackGirl_PPT_HackSRIT.pptx)")
            st.markdown("[Hackathon](https://unstop.com/hackathons/hacksrit-shri-ram-group-of-institutions-jabalpur-1471613)")
        
        st.markdown("---")
        st.caption("© 2023 Wellness Planner")

    # Main content
    tab1, tab2 = st.tabs(["Create Your Plan", "About Ayurveda"])
    
    with tab1:
        # Input Form in two columns
        st.header("Enter Your Details")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Personal Information")
            name = st.text_input("Name", placeholder="Enter your name")
            age = st.number_input("Age", min_value=1, max_value=120, value=25)
            gender = st.selectbox("Gender", options=["Male", "Female", "Other"])
            
            st.subheader("Physical Metrics")
            weight = st.number_input("Weight (kg)", min_value=1.0, max_value=300.0, value=70.0, format="%.1f")
            height = st.number_input("Height (cm)", min_value=1, max_value=250, value=170)
            
        with col2:
            st.subheader("Goals & Preferences")
            fitness_goal = st.selectbox("Fitness Goal", options=["Weight Loss", "Weight Gain", "Maintenance"])
            dietary_preference = st.selectbox("Dietary Preference", 
                                         options=["No Restrictions", "Vegetarian", "Vegan", "Keto", "Paleo", "Halal", "Gluten-Free"])
            food_allergies = st.text_input("Food Allergies (if any)", placeholder="e.g., nuts, dairy, shellfish")
            local_cuisine = st.text_input("Preferred Local Cuisine", placeholder="e.g., Indian, Italian, Chinese")
            month = st.selectbox("Current Month", 
                            options=["January", "February", "March", "April", "May", "June", 
                                     "July", "August", "September", "October", "November", "December"])
            
            include_ayurveda = st.checkbox("Include Ayurvedic wellness insights", value=True)
        
        # Calculate metrics
        bmi = round(weight / (height / 100) ** 2, 2)
        
        if bmi < 18.5:
            health_status = "Underweight"
            bmi_color = "orange"
        elif bmi <= 24.9:
            health_status = "Normal weight"
            bmi_color = "green"
        elif bmi <= 29.9:
            health_status = "Overweight"
            bmi_color = "orange"
        else:
            health_status = "Obese"
            bmi_color = "red"
            
        daily_calories = int(calculate_calorie_requirements(age, gender, weight, height, fitness_goal))
        
        # Display metrics in a nice formatted box
        st.markdown("### Your Health Metrics")
        metrics_col1, metrics_col2 = st.columns(2)
        
        with metrics_col1:
            st.markdown(f"""
            <div style="background-color: #e8f4f8; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
                <h4 style="color: #1E88E5;">BMI Analysis</h4>
                <p>Your BMI is <span style="font-weight: bold; color: {bmi_color};">{bmi}</span></p>
                <p>Status: <span style="font-weight: bold; color: {bmi_color};">{health_status}</span></p>
            </div>
            """, unsafe_allow_html=True)
            
        with metrics_col2:
            st.markdown(f"""
            <div style="background-color: #e8f4f8; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
                <h4 style="color: #1E88E5;">Calorie Requirements</h4>
                <p>Daily requirements: <span style="font-weight: bold;">{daily_calories} kcal</span></p>
                <p>Goal: <span style="font-weight: bold;">{fitness_goal}</span></p>
            </div>
            """, unsafe_allow_html=True)
        
        # User metrics dictionary
        metrics = {
            "name": name,
            "age": age,
            "gender": gender,
            "bmi": bmi,
            "health_status": health_status,
            "fitness_goal": fitness_goal,
            "dietary_preference": dietary_preference,
            "food_allergies": food_allergies,
            "daily_calories": daily_calories,
            "local_cuisine": local_cuisine,
            "weekdays": ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"],
            "month": month,
        }
        
        # Generate button
        st.markdown("---")
        generate_col1, generate_col2, generate_col3 = st.columns([1, 2, 1])
        with generate_col2:
            generate_plan = st.button("Generate Your Personalized Wellness Plan", use_container_width=True)
        
        # Generate and display plan
        if generate_plan:
            if not name:
                st.error("Please enter your name to generate a plan")
            else:
                with st.spinner("Crafting your personalized wellness plan..."):
                    try:
                        # Choose the appropriate prompt based on the Ayurveda option
                        selected_prompt = ayurveda_prompt_template if include_ayurveda else regular_prompt_template
                        plan = generate_plan_with_prompt(metrics, selected_prompt)
                        formatted_plan = format_plan(plan)
                        
                        plan_type = "Integrated Ayurvedic & Modern Wellness Plan" if include_ayurveda else "Modern Wellness Plan"
                        
                        st.success(f"Your personalized plan is ready!")
                        
                        st.markdown(f"""
                        <div class="result-container">
                            <h2 style="text-align: center; color: #2E7D32;">{plan_type} for {name}</h2>
                            <h4 style="text-align: center; color: #666;">Month: {month}</h4>
                            <hr>
                            {formatted_plan.replace('\n', '<br>')}
                        </div>
                        """, unsafe_allow_html=True)
                        
                        # Download option
                        plan_text = f"{plan_type} for {name}\nMonth: {month}\n\n{plan.content}"
                        st.download_button(
                            label="Download Your Plan",
                            data=plan_text,
                            file_name=f"{name}_wellness_plan.txt",
                            mime="text/plain"
                        )
                        
                    except Exception as e:
                        st.error(f"Error generating the plan: {e}")
    
    with tab2:
        st.header("Ayurveda: Ancient Wisdom for Modern Wellness")
        
        st.markdown("""
        <div style="background-color: white; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
            <p>Ayurveda is a traditional Indian system of medicine that dates back over 5,000 years. The term Ayurveda is derived from the Sanskrit words "ayur" (life) and "veda" (knowledge).</p>
            
            <h3>The Three Doshas</h3>
            <p>Ayurveda believes that each person has a unique mix of three energies or doshas:</p>
            <ul>
                <li><strong>Vata</strong> (Air & Space) - Controls movement and is associated with creativity and flexibility</li>
                <li><strong>Pitta</strong> (Fire & Water) - Controls metabolism and is associated with intelligence and determination</li>
                <li><strong>Kapha</strong> (Earth & Water) - Controls structure and is associated with stability and compassion</li>
            </ul>
            
            <h3>Ayurvedic Approach to Wellness</h3>
            <p>Ayurveda takes a holistic approach to health that includes:</p>
            <ul>
                <li>Personalized nutrition based on your dosha</li>
                <li>Lifestyle recommendations aligned with natural cycles</li>
                <li>Herbal medicines and therapies</li>
                <li>Mental wellness practices</li>
                <li>Seasonal adjustments to maintain balance</li>
            </ul>
            
            <p>Our wellness planner integrates these ancient principles with modern health science to provide you with a truly holistic approach to health and wellness.</p>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()