From 75be22b71905287a39ea5b2ae38af406804d1c5d Mon Sep 17 00:00:00 2001
From: Lukas Prause <lukas.prause@ikt.uni-hannover.de>
Date: Mon, 10 Jul 2023 15:53:35 +0200
Subject: [PATCH] Adds fancy legend.

---
 plot_single_transmission_EM9190.py | 35 ++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/plot_single_transmission_EM9190.py b/plot_single_transmission_EM9190.py
index 3ac3c0f..8c1ecab 100755
--- a/plot_single_transmission_EM9190.py
+++ b/plot_single_transmission_EM9190.py
@@ -200,10 +200,11 @@ if __name__ == "__main__":
         # Plot vertical lines
         first = True
         lte_handovers = transmission_df["Cell_ID"].dropna().diff()
+        lte_hanover_plot = None
         for index, value in lte_handovers.items():
             if value > 0:
                 if first:
-                    ax00.axvline(
+                    lte_hanover_plot = ax00.axvline(
                         index, ymin=0, ymax=1, color="skyblue", label="4G Handover"
                     )
                     first = False
@@ -214,29 +215,31 @@ if __name__ == "__main__":
         nr_handovers = (
             transmission_df["NR5G_Cell_ID"].replace(0, np.NaN).dropna().diff()
         )
+
+        nr_hanover_plot = None
         for index, value in nr_handovers.items():
             if value > 0:
                 if first:
-                    ax00.axvline(
+                    nr_hanover_plot = ax00.axvline(
                         index, ymin=0, ymax=1, color="greenyellow", label="5G Handover"
                     )
                     first = False
                 else:
                     ax00.axvline(index, ymin=0, ymax=1, color="greenyellow")
 
-        ax0.plot(
+        snd_plot = ax0.plot(
             transmission_df["snd_cwnd"].dropna(),
             color="lime",
             linestyle="dashed",
             label="cwnd",
         )
-        ax1.plot(
+        srtt_plot = ax1.plot(
             transmission_df["srtt"].dropna(),
             color="red",
             linestyle="dashdot",
             label="sRTT",
         )
-        ax2.plot(
+        goodput_plot = ax2.plot(
             transmission_df["goodput_rolling"],
             color="blue",
             linestyle="solid",
@@ -244,20 +247,20 @@ if __name__ == "__main__":
         )
         # ax2.plot(transmission_df["goodput"], color="blue", linestyle="solid", label="goodput")
 
-        ax01.plot(
+        eff_bw_plot = ax01.plot(
             transmission_df["effective_bw_sum"].dropna(),
             color="peru",
             linestyle="solid",
             label="bandwidth",
         )
-        ax01.plot(
+        lte_eff_bw_plot = ax01.plot(
             transmission_df["lte_effective_bw_sum"].dropna(),
             color="lightsteelblue",
             linestyle="solid",
             label="4G bandwidth",
             alpha=0.5,
         )
-        ax01.plot(
+        nr_eff_bw_plot = ax01.plot(
             transmission_df["nr_effective_bw_sum"].dropna(),
             color="cornflowerblue",
             linestyle="solid",
@@ -271,13 +274,13 @@ if __name__ == "__main__":
         #               labels=["4G bandwidth", "5G bandwidth"]
         #               )
 
-        ax02.plot(
+        lte_rsrq_plot = ax02.plot(
             transmission_df["RSRQ_(dB)"].dropna(),
             color="purple",
             linestyle="dotted",
             label="LTE RSRQ",
         )
-        ax00.plot(
+        nr_rsrq_plot = ax00.plot(
             transmission_df["NR5G_RSRQ_(dB)"].dropna(),
             color="magenta",
             linestyle="dotted",
@@ -306,7 +309,17 @@ if __name__ == "__main__":
         ax01.set_ylabel("Bandwidth [MHz]")
 
         if args.fancy:
-            fig.legend(ncols=4, fontsize=12)
+            # added these three lines
+            lns_ax0 = snd_plot + srtt_plot + goodput_plot
+            labs_ax0 = [l.get_label() for l in lns_ax0]
+            ax0.legend(lns_ax0, labs_ax0, ncols=4, fontsize=12, loc="upper center")
+            lns_ax00 = eff_bw_plot + lte_eff_bw_plot + nr_eff_bw_plot + lte_rsrq_plot + nr_rsrq_plot
+            if lte_hanover_plot:
+                lns_ax00.append(lte_hanover_plot)
+            if nr_hanover_plot:
+                lns_ax00.append(nr_hanover_plot)
+            labs_ax00 = [l.get_label() for l in lns_ax0]
+            ax00.legend(lns_ax00, labs_ax00, ncols=4, fontsize=12, loc="upper center")
             plt.savefig("{}{}_plot.eps".format(args.save, csv.replace(".csv", "")), bbox_inches="tight")
         else:
             fig.legend(loc="lower right")