Update app.py

app.py

@@ -5,43 +5,31 @@ from model_api import query_model
 # 1. Page Configuration
 st.set_page_config(page_title="FitPlan AI", layout="wide")
 
-# 2. Custom Styling (Preserving Sidebar and Black Text)
+# 2. Custom Styling
 st.markdown("""
     <style>
-    .stApp { 
-        background: linear-gradient(135deg, #f6d365 0%, #fda085 100%); 
-        background-attachment: fixed; 
-    }
-    section[data-testid="stSidebar"] {
-        background-color: #000000 !important;
-    }
-    section[data-testid="stSidebar"] h1, section[data-testid="stSidebar"] p {
-        color: #FFFFFF !important;
-    }
-    h1, h2, h3, p, label { 
-        color: #2D3436 !important; 
-        font-weight: 600; 
-    }
+    .stApp { background: linear-gradient(135deg, #f6d365 0%, #fda085 100%); background-attachment: fixed; }
+    section[data-testid="stSidebar"] { background-color: #000000 !important; }
+    section[data-testid="stSidebar"] h1, section[data-testid="stSidebar"] p { color: #FFFFFF !important; }
+    h1, h2, h3, p, label { color: #2D3436 !important; font-weight: 600; }
     .result-card { 
-        background-color: white !important; 
-        padding: 25px; 
-        border-radius: 15px; 
-        border-left: 10px solid #004d57; 
-        color: #000000 !important; 
-        margin-bottom: 20px;
+        background-color: white !important; padding: 25px; border-radius: 15px; 
+        border-left: 10px solid #004d57; color: #000000 !important; margin-bottom: 20px;
     }
     .day-card {
-        background-color: #ffffff !important; 
-        padding: 20px; 
-        border-radius: 12px;
-        margin-bottom: 15px; 
-        border-left: 8px solid #004d57; 
-        color: #000000 !important;
+        background-color: #ffffff !important; padding: 20px; border-radius: 12px;
+        margin-bottom: 15px; border-left: 8px solid #004d57; color: #000000 !important;
         white-space: pre-wrap !important;
     }
     </style>
     """, unsafe_allow_html=True)
 
+# Initialize Session State variables if they don't exist
+if 'workout_plan' not in st.session_state:
+    st.session_state.workout_plan = None
+if 'assessment_html' not in st.session_state:
+    st.session_state.assessment_html = None
+
 # --- SIDEBAR ---
 with st.sidebar:
     st.markdown("# 👤 Personalised Fitness Plan")
@@ -51,70 +39,63 @@ with st.sidebar:
 # --- MAIN CONTENT ---
 st.title("🏋️ Create Your Fitness Profile")
 
-# Input Section
 col1, col2, col3 = st.columns([2, 1, 1])
-with col1: 
-    name = st.text_input("Name")
-with col2: 
-    gender = st.selectbox("Gender", ["Male", "Female", "Other"])
-with col3: 
-    age = st.number_input("Age", min_value=1, value=19)
+with col1: name = st.text_input("Name")
+with col2: gender = st.selectbox("Gender", ["Male", "Female", "Other"])
+with col3: age = st.number_input("Age", min_value=1, value=19)
 
 col_h, col_w = st.columns(2)
-with col_h: 
-    height = st.number_input("Height (cm)", min_value=1.0)
-with col_w: 
-    weight = st.number_input("Weight (kg)", min_value=1.0)
+with col_h: height = st.number_input("Height (cm)", min_value=1.0)
+with col_w: weight = st.number_input("Weight (kg)", min_value=1.0)
 
 goal = st.selectbox("Goal", ["Build Muscle", "Weight Loss", "Strength", "Flexibility"])
 level = st.selectbox("Fitness Level", ["Beginner", "Intermediate", "Advanced"])
 equip = st.multiselect("Equipment", ["Dumbbells", "Kettlebells", "Pull-up Bar", "Resistance Bands", "Yoga Mat", "No Equipment"])
 
-# --- SUBMIT BUTTON LOGIC ---
+# --- SUBMIT BUTTON ---
 if st.button("Submit Profile"):
     if not name or height <= 0 or weight <= 0:
         st.error("Please fill in all details correctly.")
     else:
-        # 1. Build the prompt (Expects 4 values from prompt_builder.py)
         prompt, bmi, bmi_status, status_color = build_prompt(name, gender, height, weight, goal, level, equip)
 
-        # 2. Display BMI Assessment Card
-        st.markdown(f"""
+        # Save assessment to session state
+        st.session_state.assessment_html = f"""
             <div class="result-card">
                 <h2>Assessment for {name}</h2>
                 <p>Gender: {gender} | Age: {age} Years</p>
                 <p>Calculated BMI: <strong>{bmi:.2f}</strong> | Category: <span style="color:{status_color}">{bmi_status}</span></p>
             </div>
-        """, unsafe_allow_html=True)
+        """
 
-        # 3. Call Inference API
         with st.spinner("AI Trainer is thinking..."):
             raw_workout = query_model(prompt)
-
             if "Error" in raw_workout:
                 st.error(f"API Error: {raw_workout}")
             else:
-                st.markdown("### 📋 Your 5-Day Routine")
-                
-                # Split and Display Days
-                days = ["Day 1:", "Day 2:", "Day 3:", "Day 4:", "Day 5:"]
-                
-                # We iterate through the days to create individual white cards
-                for i in range(len(days)):
-                    start = raw_workout.find(days[i])
-                    if start != -1:
-                        # Find where the next day starts to slice the text
-                        end = raw_workout.find(days[i+1]) if i+1 < len(days) else len(raw_workout)
-                        day_content = raw_workout[start:end].strip()
-                        
-                        # Render the card with black text
-                        st.markdown(f'<div class="day-card">{day_content}</div>', unsafe_allow_html=True)
-                
-                # --- DOWNLOAD BUTTON (Correctly indented inside the block) ---
-                st.write("---")
-                st.download_button(
-                    label="📥 Download Full Plan",
-                    data=raw_workout,
-                    file_name=f"{name}_Workout_Plan.txt",
-                    mime="text/plain"
-                )
+                # Save workout to session state so it persists during reload
+                st.session_state.workout_plan = raw_workout
+
+# --- DISPLAY SECTION (Outside the button to handle reloads) ---
+if st.session_state.assessment_html:
+    st.markdown(st.session_state.assessment_html, unsafe_allow_html=True)
+
+if st.session_state.workout_plan:
+    st.markdown("### 📋 Your 5-Day Routine")
+    raw_workout = st.session_state.workout_plan
+    days = ["Day 1:", "Day 2:", "Day 3:", "Day 4:", "Day 5:"]
+    
+    for i in range(len(days)):
+        start = raw_workout.find(days[i])
+        if start != -1:
+            end = raw_workout.find(days[i+1]) if i+1 < len(days) else len(raw_workout)
+            day_content = raw_workout[start:end].strip()
+            st.markdown(f'<div class="day-card">{day_content}</div>', unsafe_allow_html=True)
+    
+    st.write("---")
+    st.download_button(
+        label="📥 Download Full Plan",
+        data=st.session_state.workout_plan,
+        file_name=f"{name}_Workout_Plan.txt",
+        mime="text/plain"
+    )