From 9bad49dbec9200f4a0e769e4505105c807b33fb0 Mon Sep 17 00:00:00 2001
From: Merlin Unterfinger <merlin.unterfinger@sbb.ch>
Date: Mon, 7 Oct 2024 16:38:23 +0200
Subject: [PATCH] enh: add overview plot to depots

---
 rssched/app/pages/04_depots.py | 13 +++++----
 rssched/app/utils/plot.py      | 48 ++++++++++++++++++++++++++++++++++
 rssched/app/utils/transform.py |  4 ++-
 3 files changed, 59 insertions(+), 6 deletions(-)
 create mode 100644 rssched/app/utils/plot.py

diff --git a/rssched/app/pages/04_depots.py b/rssched/app/pages/04_depots.py
index c2bb14b..651f69e 100644
--- a/rssched/app/pages/04_depots.py
+++ b/rssched/app/pages/04_depots.py
@@ -1,9 +1,7 @@
-from typing import List
-
-import pandas as pd
 import streamlit as st
 
 from rssched.app.utils.io import get_uploaded_data
+from rssched.app.utils.plot import plot_depots_bar_chart
 from rssched.app.utils.transform import flatten_depots
 from rssched.visualization.depot_loads import plot_depot_vehicle_loads
 
@@ -17,11 +15,16 @@
 
 
 with tabs[0]:
+    st.plotly_chart(plot_depots_bar_chart(df_depots))
+
     show_aggregated = st.checkbox("Aggregate", value=True)
 
     if show_aggregated:
-        df_grouped = df_depots.groupby("depot_id", as_index=False).agg(
-            {"vehicles": "sum", "capacity": "mean"}
+        df_grouped = (
+            df_depots.groupby("depot_id", as_index=False)
+            .agg({"vehicles": "sum", "capacity": "mean"})
+            .sort_values(by="vehicles")
+            .reset_index(drop=True)
         )
         st.dataframe(df_grouped, hide_index=True)
     else:
diff --git a/rssched/app/utils/plot.py b/rssched/app/utils/plot.py
new file mode 100644
index 0000000..e3210f9
--- /dev/null
+++ b/rssched/app/utils/plot.py
@@ -0,0 +1,48 @@
+import pandas as pd
+import plotly.graph_objects as go
+from plotly.graph_objs import Figure
+
+
+def plot_depots_bar_chart(df_depots: pd.DataFrame) -> Figure:
+    # Group by depot_id to calculate the total capacity and sort by the number of vehicles
+    df_depots_grouped = (
+        df_depots[df_depots["vehicles"] > 0]
+        .groupby("depot_id", as_index=False)
+        .agg({"vehicles": "sum", "capacity": "max"})
+        .sort_values(by="vehicles", ascending=False)
+    )
+
+    # Create a bar chart
+    fig = go.Figure()
+
+    # Add bar for the number of vehicles
+    fig.add_trace(
+        go.Bar(
+            x=df_depots_grouped["depot_id"],
+            y=df_depots_grouped["vehicles"],
+            name="Vehicles",
+            marker_color="blue",  # Color for vehicles bar
+        )
+    )
+
+    # Add bar for the depot capacity
+    fig.add_trace(
+        go.Bar(
+            x=df_depots_grouped["depot_id"],
+            y=df_depots_grouped["capacity"],
+            name="Capacity",
+            marker_color="gray",  # Color for capacity bar
+        )
+    )
+
+    # Update layout of the bar chart
+    fig.update_layout(
+        barmode="group",  # Bars grouped side by side
+        title="Depot Vehicles and Capacity",
+        xaxis_title="Depot ID",
+        # yaxis_title="Count",
+        legend_title="Legend",
+    )
+
+    # Return the figure object
+    return fig
diff --git a/rssched/app/utils/transform.py b/rssched/app/utils/transform.py
index 33ee4fb..9e9ebd3 100644
--- a/rssched/app/utils/transform.py
+++ b/rssched/app/utils/transform.py
@@ -14,7 +14,9 @@ def flatten_depots(request: Request, response: Response) -> pd.DataFrame:
         df_request, df_response, on=["depot_id", "vehicle_type"], how="outer"
     )
     df_combined["vehicles"] = df_combined["vehicles"].fillna(0).astype(int)
-    return df_combined
+    return df_combined.sort_values(by="vehicles", ascending=False).reset_index(
+        drop=True
+    )
 
 
 def _flatten_response_depot_loads(