Optimise get workout counts function (2-3X faster however it was already pretty fast, need to find more slow shit)

utils.py

@@ -309,39 +309,37 @@ def get_exercise_graph_model(title, estimated_1rm, repetitions, weight, start_da
     }
 
 def get_workout_counts(workouts, period='week'):
-    # Convert to DataFrame
     df = pd.DataFrame(workouts)
     
-    # Convert 'StartDate' to datetime
+    # Convert 'StartDate' to datetime and set period
     df['StartDate'] = pd.to_datetime(df['StartDate'])
+    df['Period'] = df['StartDate'].dt.to_period('W' if period == 'week' else 'M')
 
-    # Determine the range of periods to cover
-    min_date = df['StartDate'].min()
-    max_date = pd.Timestamp(datetime.now())
+    # Group by PersonId, Period and count unique workouts
+    workout_counts = df.groupby(['PersonId', 'Period'])['WorkoutId'].nunique().reset_index()
 
-    # Generate a complete range of periods
-    freq = 'W-MON' if period == 'week' else 'MS'
-    period_range = pd.date_range(start=min_date, end=max_date, freq=freq)
+    # Convert 'Period' to timestamp using the start date of the period
+    workout_counts['Period'] = workout_counts['Period'].apply(lambda x: x.start_time)
 
-    # Initialize a dictionary to store workout counts and person names
-    workout_counts = {
-        person_id: {
-            "PersonName": person_name,
-            "PRCounts": {p: 0 for p in period_range}
-        } for person_id, person_name in df[['PersonId', 'PersonName']].drop_duplicates().values
-    }
+    # Pivot the result to get periods as columns
+    workout_counts_pivot = workout_counts.pivot(index='PersonId', columns='Period', values='WorkoutId').fillna(0)
 
-    # Process the workouts
-    for person_id, person_data in workout_counts.items():
-        person_df = df[df['PersonId'] == person_id]
+    # Include person names
+    names = df[['PersonId', 'PersonName']].drop_duplicates().set_index('PersonId')
+    workout_counts_final = names.join(workout_counts_pivot, how='left').fillna(0)
 
-        for period_start in person_data["PRCounts"]:
-            period_end = period_start + pd.DateOffset(weeks=1) if period == 'week' else period_start + pd.DateOffset(months=1)
-            period_workouts = person_df[(person_df['StartDate'] >= period_start) & (person_df['StartDate'] < period_end)]
-            period_workout_count = period_workouts['WorkoutId'].unique()
-            person_data["PRCounts"][period_start] = len(period_workout_count)
+    # Convert DataFrame to dictionary
+    result = workout_counts_final.reset_index().to_dict('records')
 
-    return workout_counts
+    # Reformat the dictionary to desired structure
+    formatted_result = {}
+    for record in result:
+        person_id = record.pop('PersonId')
+        person_name = record.pop('PersonName')
+        pr_counts = {k: v for k, v in record.items()}
+        formatted_result[person_id] = {'PersonName': person_name, 'PRCounts': pr_counts}
+
+    return formatted_result
 
 def count_prs_over_time(workouts, period='week'):
     df = pd.DataFrame(workouts)