From 11a4cf42485667aa0ecf8927b7b17991f6ee55cb Mon Sep 17 00:00:00 2001
From: Felix Blanke <felixblanke@uni-bonn.de>
Date: Mon, 28 Aug 2023 17:08:03 +0200
Subject: [PATCH] Allow filtering for landesbezirk

---
 wsgi.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/wsgi.py b/wsgi.py
index 8e6ed3f..cb002f1 100644
--- a/wsgi.py
+++ b/wsgi.py
@@ -90,6 +90,7 @@ def plot(
     sheet_name: str = "digital",
     total_targets: tuple[int, ...] = (1500, ),
     alpha: float | None = None,
+    landesbez_str: str | None = None
 ) -> str:
     curr_datetime = datetime.datetime.now()
     df = create_plot_df(
@@ -108,33 +109,39 @@ def plot(
         if day.weekday() >= 5:
             plt.gca().axvspan(days[idx], days[idx + 1], alpha=0.2, color="gray")
 
+    series = df.sum(axis=1) if landesbez_str is None else df[landesbez_str]
+    plot_df = series.to_frame("Digitale Befragung").replace(0, np.nan)
+    plot_df = plot_df.astype({"Digitale Befragung": "float32"})
+    if pd.isna(plot_df).all().item():
+        continue
     if alpha is not None:
         plt.fill_between(
-            df.dropna().index,
-            df.dropna()["Digitale Befragung"],
+            plot_df.dropna().index,
+            plot_df.dropna()["Digitale Befragung"],
             color="#e4004e",
             alpha=alpha,
         )
 
     plt.plot(
-        df.dropna().index,
-        df.dropna()["Digitale Befragung"],
+        plot_df.dropna().index,
+        plot_df.dropna()["Digitale Befragung"],
         ls="--",
         marker="o",
         lw=1,
         color="#e4004e",
         markersize=4,
+        label=landesbez_str,
     )
 
     if current_df is not None:
         plt.annotate(
             "Jetzt",
-            (df.dropna().index[-1], df.dropna()["Digitale Befragung"][-1] * 1.03),
+            (plot_df.dropna().index[-1], plot_df.dropna()["Digitale Befragung"][-1] * 1.03),
             fontsize=8,
             ha="center",
         )
 
-    plt.plot(df.index, df["Digitale Befragung"], lw=1.5, color="#e4004e")
+    plt.plot(plot_df.index, plot_df["Digitale Befragung"], lw=1.5, color="#e4004e", label=landesbez_str,)
 
     plt.title("Teilnahme an Digitaler Beschäftigtenbefragung")
     plt.ylabel("# Teilnahmen")