Update app.py

app.py

@@ -137,21 +137,11 @@ with col3:
 
 
 
-# Ensure the session state variables are initialized
-if 'current_page' not in st.session_state:
-    st.session_state.current_page = 'home'
-if 'premiums_df' not in st.session_state:
-    st.session_state.premiums_df = None
-if 'claims_recoveries_df' not in st.session_state:
-    st.session_state.claims_recoveries_df = None
-if 'claims_count_df' not in st.session_state:
-    st.session_state.claims_count_df = None
-
-if 'current_page' not in st.session_state:
-    st.session_state['current_page'] = 'home'
-if 'premiums_df' not in st.session_state:
-    st.session_state['premiums_df'] = None
-# ... Similarly for other dataframes
+# Initialize session state variables
+st.session_state.setdefault('current_page', 'home')
+st.session_state.setdefault('premiums_df', None)
+st.session_state.setdefault('claims_recoveries_df', None)
+st.session_state.setdefault('claims_count_df', None)
 
 # Helper function to check if a string has more than 50% numeric characters
 def is_mostly_numeric(s):
@@ -180,7 +170,7 @@ def is_mostly_numeric(s):
 def process_file(uploaded_file, key):
     df = pd.read_csv(uploaded_file)
     st.session_state[key] = df
-    st.session_state['current_page'] = key
+    st.session_state['current_page'] = key + '_selection'
 
 # Function to display the column selection UI
 def column_selection_ui(dataframe_key, selection_key):
@@ -223,42 +213,45 @@ def column_selection_ui(dataframe_key, selection_key):
             # ... include other selections as well
         }
 
+    # Store the column selections in the session state
+    st.session_state[f'{selection_key}_selections'] = {
+        'lob_column': lob_column,
+        'selected_lob_values': selected_lob_values,
+        'valuation_as_at': valuation_as_at,
+        'quarter_bracket': quarter_bracket,
+    }
 
-# File uploader and buttons
-upload_col1, upload_col2, upload_col3 = st.columns(3)
-with upload_col1:
-    uploaded_file_premiums = st.file_uploader("Upload Premiums Data", key="premiums")
-with upload_col2:
-    uploaded_file_claims_recoveries = st.file_uploader("Upload Claims Recoveries Data", key="claims_recoveries")
-with upload_col3:
-    uploaded_file_claims_count = st.file_uploader("Upload Claims Count Data", key="claims_count")
+    
+# File uploaders and buttons
+if st.session_state.current_page == 'home':
+    upload_col1, upload_col2, upload_col3 = st.columns(3)
+    with upload_col1:
+        uploaded_file_premiums = st.file_uploader("Upload Premiums Data", key="premiums")
+    with upload_col2:
+        uploaded_file_claims_recoveries = st.file_uploader("Upload Claims Recoveries Data", key="claims_recoveries")
+    with upload_col3:
+        uploaded_file_claims_count = st.file_uploader("Upload Claims Count Data", key="claims_count")
 
-# Only show buttons if we're on the home page
-if st.session_state['current_page'] == 'home':
     button_col1, button_col2, button_col3 = st.columns(3)
     with button_col1:
         if uploaded_file_premiums and st.button("Select Premiums columns"):
-            process_file(uploaded_file_premiums, 'premiums_df')
+            process_file(uploaded_file_premiums, 'premiums')
     with button_col2:
         if uploaded_file_claims_recoveries:
             if st.button("Select Incurred claims columns"):
-                process_file(uploaded_file_claims_recoveries, 'claims_recoveries_df')
+                process_file(uploaded_file_claims_recoveries, 'claims_recoveries')
             if st.button("Select Recoveries columns"):
-                # Use a different key or logic for processing recoveries
-                process_file(uploaded_file_claims_recoveries, 'claims_recoveries_df')
+                process_file(uploaded_file_claims_recoveries, 'claims_recoveries')
     with button_col3:
         if uploaded_file_claims_count and st.button("Select Claims count columns"):
-            process_file(uploaded_file_claims_count, 'claims_count_df')
-
-# Display the column selection UI based on the current page
-if st.session_state['current_page'] == 'premiums_df':
-    column_selection_ui('premiums_df', 'premiums')
-elif st.session_state['current_page'] == 'claims_recoveries_df':
-    column_selection_ui('claims_recoveries_df', 'claims_recoveries')
-elif st.session_state['current_page'] == 'claims_count_df':
-    column_selection_ui('claims_count_df', 'claims_count')
-
-# Add a button to go back to the home page from any selection page
-if st.session_state['current_page'] != 'home':
+            process_file(uploaded_file_claims_count, 'claims_count')
+
+# Conditional rendering based on the current page
+if st.session_state.current_page.endswith('_selection'):
+    dataframe_key = st.session_state.current_page.replace('_selection', '')
+    column_selection_ui(dataframe_key, dataframe_key)
+
+# Back to Home button
+if st.session_state.current_page != 'home':
     if st.button("Back to Home"):
-        st.session_state['current_page'] = 'home'
+        st.session_state.current_page = 'home'