app.py

import streamlit as st
from dataclasses import dataclass
from typing import List, Dict, Optional
from datetime import datetime
import google.generativeai as genai
import os
from dotenv import load_dotenv
from PIL import Image
from langchain_groq import ChatGroq
from langchain.schema import HumanMessage, SystemMessage
import torch
import numpy as np
import plotly.graph_objects as go
from fpdf import FPDF
import time

# Load environment variables
load_dotenv()

# Configure LLMs
llm = ChatGroq(
    temperature=0.2,
    groq_api_key=os.environ.get("GROQ_API_KEY"),
    model_name="llama-3.1-70b-versatile"
)

genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config={
        "temperature": 0.2,
        "top_p": 0.8,
        "top_k": 40,
        "max_output_tokens": 2048,
    }
)

# Initialize session state for realtime.py components
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'user_data' not in st.session_state:
    st.session_state.user_data = {
        'progress': [],
        'meal_logs': [],
        'workout_logs': []
    }

# Code from acountable.py
@dataclass
class TransformationJourney:
    user_id: str
    goal_type: str  # fitness, skill, personal growth etc
    start_date: datetime
    progress_images: List[Dict[str, str]]  # [{date: image_path}]
    feedback_history: List[Dict[str, str]]  # [{date: feedback}]
    chat_history: List[Dict[str, str]]  # [{role: content}]

class TransformationCoach:
    def __init__(self):
        self.model = model
        self.llm = llm
        
    def analyze_image(self, image: Image) -> str:
        prompt = """
        As a transformation coach, analyze this image in relation to the user's goals.
        Provide detailed observations about:
        1. Current state/progress visible in the image
        2. Areas of improvement
        3. Notable achievements
        Be specific and objective in your analysis.
        """
        response = self.model.generate_content([prompt, image])
        return response.text
    
    def generate_feedback(self, analysis: str, goal_type: str) -> Dict[str, List[str]]:
        prompt = f"""
        Based on the following analysis of the user's progress image for their {goal_type} transformation:
        {analysis}
        
        Provide structured feedback in the following format:
        1. Specific "DO" recommendations (3-5 points)
        2. Specific "DON'T" recommendations (2-3 points)
        3. A motivational message
        Make all feedback actionable and tailored to the analysis.
        """
        response = self.model.generate_content(prompt)
        feedback_parts = response.text.split("\n\n")
        return {
            "dos": feedback_parts[0].split("\n")[1:],
            "donts": feedback_parts[1].split("\n")[1:],
            "motivation": feedback_parts[2]
        }

    def chat_response(self, question: str, context: List[Dict[str, str]], analysis: str = None) -> str:
        prompt = f"""
        You are a supportive transformation coach. Based on the following context:
        
        Previous chat: {str(context)}
        Latest analysis: {analysis if analysis else 'No recent analysis'}
        
        User question: {question}
        
        Provide a helpful, encouraging response focused on their transformation journey.
        """
        response = self.llm.invoke(prompt)
        return response.content

# Functions from realtime.py
def load_llm():
    api_key = os.environ.get("GROQ_API_KEY")
    if api_key is None:
        st.error("GROQ_API_KEY is not set in the environment variables.")
        return None
    llm = ChatGroq(
        api_key=api_key,
        model_name="llama-3.1-70b-versatile",
        temperature=0.7,
        max_tokens=2048
    )
    return llm

def analyze_image_realtime(image, task="meal"):
    try:
        if task == "meal":
            prompt = """
            As a nutrition expert, analyze this meal image and provide:
            1. Estimated caloric content
            2. Macro-nutrient breakdown
            3. Nutritional assessment
            4. Suggestions for improvement
            5. Any potential health concerns
            Be specific and detailed in your analysis.
            """
        else:  # workout form analysis
            prompt = """
            As a fitness expert, analyze this workout form and provide:
            1. Form assessment
            2. Potential injury risks
            3. Specific corrections needed
            4. Benefits of proper form
            Be specific and detailed in your analysis.
            """
        response = model.generate_content([prompt, image])
        return response.text
    except Exception as e:
        return f"Error analyzing image: {str(e)}"

def track_progress(metric, value, date=None):
    if date is None:
        date = datetime.now().strftime("%Y-%m-%d")
    
    st.session_state.user_data['progress'].append({
        'date': date,
        'metric': metric,
        'value': value
    })

def display_progress_chart():
    if not st.session_state.user_data['progress']:
        return
    
    data = st.session_state.user_data['progress']
    fig = go.Figure()
    
    metrics = set(item['metric'] for item in data)
    for metric in metrics:
        metric_data = [item for item in data if item['metric'] == metric]
        dates = [item['date'] for item in metric_data]
        values = [item['value'] for item in metric_data]
        
        fig.add_trace(go.Scatter(x=dates, y=values, name=metric))
    
    fig.update_layout(title="Your Progress Over Time",
                     xaxis_title="Date",
                     yaxis_title="Value")
    st.plotly_chart(fig)

# Functions from workout.py
def generate_meal_plan(weight, height, age, gender, activity_level, goal, dietary_restrictions, health_history, current_eating_habits, meal_preparation, budget, sleep):
    """Generate personalized meal plan using LLM"""
    prompt = f"""
    Based on the following user data, generate a highly personalized and detailed daily meal plan with easily available ingredients:
    
    User Profile:
    - Weight: {weight} kg
    - Height: {height} cm 
    - Age: {age}
    - Gender: {gender}
    - Activity Level: {activity_level}
    - Goal: {goal}
    - Dietary Restrictions: {', '.join(dietary_restrictions)}
    - Health History: {health_history}
    - Current Eating Habits: {current_eating_habits}
    - Meal Preparation: {meal_preparation}
    - Budget: {budget}
    - Sleep: {sleep}

    Generate a clear and structured meal plan with:

    1. Total Daily Calorie Requirements
    2. Macronutrient Distribution (in grams):
       - Protein
       - Carbohydrates  
       - Fats

    Daily Meal Schedule:
    [Detailed meal schedule content...]
    """
    
    response = llm.invoke(prompt)
    return response.content

def generate_workout_plan(goal, activity_level, target_areas, workout_time, current_fitness_level, exercise_preferences):
    """Generate personalized workout plan using LLM"""
    prompt = f"""
    Based on:
    - Fitness Goal: {goal}
    - Activity Level: {activity_level}
    - Target Areas: {target_areas}
    - Workout Time: {workout_time}
    - Current Fitness Level: {current_fitness_level}
    - Exercise Preferences: {exercise_preferences}

    Generate a clear and structured weekly workout plan:
    [Detailed workout plan content...]
    """
    
    response = llm.invoke(prompt)
    return response.content

def create_pdf(name, meal_plan, workout_plan):
    pdf = FPDF()
    pdf.add_page()
    
    pdf.set_font('Arial', 'B', 20)
    pdf.cell(0, 10, f'Personalized Fitness Plan for {name}', ln=True, align='C')
    pdf.ln(10)
    
    pdf.set_font('Arial', 'B', 16)
    pdf.cell(0, 10, 'Meal Plan', ln=True)
    pdf.set_font('Arial', '', 12)
    pdf.multi_cell(0, 10, meal_plan)
    
    pdf.add_page()
    pdf.set_font('Arial', 'B', 16)
    pdf.cell(0, 10, 'Workout Plan', ln=True)
    pdf.set_font('Arial', '', 12)
    pdf.multi_cell(0, 10, workout_plan)
    
    return pdf

def main():
    st.set_page_config(page_title="AI Fitness Coach", layout="wide")
    
    # Chat input needs to be outside tabs
    user_input = st.chat_input("Type your message here...")
    if user_input:
        st.session_state.chat_history.append({"role": "user", "content": user_input})
        
        # Process user input with LLM
        llm = load_llm()
        if llm is not None:
            messages = [
                SystemMessage(content="You are an expert health and fitness coach. Provide detailed and personalized advice."),
                HumanMessage(content=user_input)
            ]
            response = llm(messages)
            
            st.session_state.chat_history.append({"role": "assistant", "content": response.content})
    
    tab1, tab2, tab3 = st.tabs(["Transformation Journey", "Realtime Tracking", "Workout & Meal Plans"])
    
    with tab1:
        # Implementation from acountable.py
        st.title("Transformation Journey Tracker")
        
        if 'coach' not in st.session_state:
            st.session_state.coach = TransformationCoach()
        if 'journey' not in st.session_state:
            st.session_state.journey = None
            
        # Initialize or load journey
        if not st.session_state.journey:
            with st.form("journey_setup"):
                st.write("Let's start your transformation journey!")
                user_id = st.text_input("Enter your user ID")
                goal_type = st.selectbox(
                    "What type of transformation?",
                    ["Fitness", "Skill Development", "Personal Growth", "Other"]
                )
                if st.form_submit_button("Start Journey"):
                    st.session_state.journey = TransformationJourney(
                        user_id=user_id,
                        goal_type=goal_type,
                        start_date=datetime.now(),
                        progress_images=[],
                        feedback_history=[],
                        chat_history=[]
                    )
        
        # Main interaction area
        if st.session_state.journey:
            st.subheader("Upload Progress Image")
            uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"], key="transform_uploader")
            
            if uploaded_file:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded Image", use_column_width=True)
                
                if st.button("Analyze and Get Feedback"):
                    with st.spinner("Analyzing your progress..."):
                        analysis = st.session_state.coach.analyze_image(image)
                        feedback = st.session_state.coach.generate_feedback(
                            analysis, 
                            st.session_state.journey.goal_type
                        )
                        
                        # Display feedback
                        st.subheader("Your Personalized Feedback")
                        
                        st.write("🎯 DO:")
                        for do in feedback["dos"]:
                            st.write(f"✓ {do}")
                            
                        st.write("⚠️ DON'T:")
                        for dont in feedback["donts"]:
                            st.write(f"✗ {dont}")
                            
                        st.write("💪 Motivation:")
                        st.write(feedback["motivation"])
                        
                        # Store progress
                        current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
                        image_path = f"progress_images/{st.session_state.journey.user_id}_{current_time}.jpg"
                        st.session_state.journey.progress_images.append({
                            current_time: image_path
                        })
                        st.session_state.journey.feedback_history.append({
                            current_time: feedback
                        })
            
            # Chat interface
            st.write("---")
            st.subheader("Chat with Your Coach")
            
            # Display chat history
            for message in st.session_state.journey.chat_history:
                if message["role"] == "user":
                    st.write("You: " + message["content"])
                else:
                    st.write("Coach: " + message["content"])
            
            # Chat input
            user_question = st.text_input("Ask your coach anything:", key="transform_chat")
            if st.button("Send", key="transform_send"):
                if user_question:
                    # Add user message
                    st.session_state.journey.chat_history.append({
                        "role": "user",
                        "content": user_question
                    })
                    
                    # Get coach's response
                    latest_analysis = st.session_state.coach.analyze_image(image) if uploaded_file else None
                    response = st.session_state.coach.chat_response(
                        user_question,
                        st.session_state.journey.chat_history,
                        latest_analysis
                    )
                    
                    # Add coach's response
                    st.session_state.journey.chat_history.append({
                        "role": "coach",
                        "content": response
                    })
                    
                    st.experimental_rerun()

    with tab2:
        # Implementation from realtime.py
        st.title("AI Health Coach")
        
        st.warning("""This system is for informational purposes only and is not a substitute for professional medical advice, 
                    diagnosis, or treatment. Always seek the advice of your physician or other qualified health provider.""")
        
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.write(message["content"])
        
        # Image upload and analysis
        uploaded_file = st.file_uploader("Upload an image for analysis (meal or workout form)", type=['png', 'jpg', 'jpeg'], key="realtime_uploader")
        if uploaded_file:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image")
            
            analysis_type = st.radio("What would you like to analyze?", ["Meal", "Workout Form"])
            if st.button("Analyze Image", key="realtime_analyze"):
                with st.spinner("Analyzing image..."):
                    results = analyze_image_realtime(image, analysis_type.lower())
                    st.write("Analysis Results:", results)
        
        # Progress tracking section
        with st.expander("Track Your Progress"):
            metric = st.selectbox("Select metric to track", ["Weight", "Body Fat %", "Workout Duration", "Calories"])
            value = st.number_input("Enter value", key="progress_value")
            if st.button("Log Progress", key="log_progress"):
                track_progress(metric, value)
                st.success("Progress logged successfully!")
        
        # Display progress charts
        display_progress_chart()

    with tab3:
        # Implementation from workout.py
        st.title("AI Fitness Coach")
        st.markdown("<h1 style='text-align: center; color: #4CAF50;'>Welcome to Your Personalized Fitness Journey!</h1>", unsafe_allow_html=True)
        
        # Get user inputs
        st.header("Personal Information")

        col1, col2, col3 = st.columns(3)

        with col1:
            name = st.text_input("What's your name?", placeholder="Enter your name")
            age = st.number_input("What's your age?", min_value=1, max_value=120, value=25)
            gender = st.selectbox("What's your gender?", ["Male", "Female"])
            health_history = st.text_input("Do you have any medical conditions?")

        with col2:
            weight = st.number_input("What's your weight (in kg)?", min_value=1, value=70)
            height = st.number_input("What's your height (in cm)?", min_value=1, value=170)
            goal = st.selectbox(
                "What's your fitness goal?",
                ["Weight Loss", "Muscle Gain", "Maintenance", "General Fitness"],
                index=0
            )
            activity_level = st.selectbox(
                "What's your activity level?",
                [
                    "Sedentary (little to no exercise)",
                    "Lightly Active (1-3 days/week)",
                    "Moderately Active (3-5 days/week)",
                    "Very Active (6-7 days/week)",
                    "Extra Active (very intense exercise daily)"
                ]
            )

        with col3:
            dietary_restrictions = st.multiselect(
                "Do you have any dietary restrictions?",
                ["None", "Vegetarian", "Vegan", "Gluten-Free", "Dairy-Free", "Keto"]
            )
            current_eating_habits = st.text_input("What does your current diet look like?")
            meal_preparation = st.text_input("Do you cook your meals or rely on takeout?")
            budget = st.number_input("What is your monthly budget for groceries/food?", min_value=0)

        if st.button("Generate Plans", key="generate_plans"):
            if not name or not goal or not age or not gender or not weight or not height or not activity_level:
                st.error("Please fill in all required fields")
                return
                
            with st.spinner("Generating your personalized plans..."):
                meal_plan = generate_meal_plan(
                    weight=weight,
                    height=height,
                    age=age,
                    gender=gender,
                    activity_level=activity_level,
                    goal=goal,
                    dietary_restrictions=dietary_restrictions if dietary_restrictions else ["None"],
                    health_history=health_history,
                    current_eating_habits=current_eating_habits,
                    meal_preparation=meal_preparation,
                    budget=budget,
                    sleep="8 hours"  # Default value
                )
                workout_plan = generate_workout_plan(
                    goal=goal,
                    activity_level=activity_level,
                    target_areas="full body",  # Default value
                    workout_time=60,  # Default value
                    current_fitness_level="Beginner",  # Default value
                    exercise_preferences="None"  # Default value
                )
                
                # Create PDF
                pdf = create_pdf(name, meal_plan, workout_plan)
                
                # Display plans in two columns
                col1, col2 = st.columns(2)
                with col1:
                    st.subheader("Your Meal Plan")
                    st.write(meal_plan)
                with col2:
                    st.subheader("Your Workout Plan")
                    st.write(workout_plan)
                
                # Add download button
                st.download_button(
                    label="Download Plan as PDF",
                    data=pdf.output(dest='S').encode('latin-1'),
                    file_name=f"fitness_plan_{name}.pdf",
                    mime="application/pdf"
                )

if __name__ == "__main__":
    main()