From c231fd346078c0fe6a9eb6d686246b9fbe7e6c71 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Tue, 14 May 2024 14:03:48 +0900
Subject: [PATCH] add plot for message sizes

---
 mixnet/v2/sim/main.py | 13 +++++++++++++
 mixnet/v2/sim/p2p.py  |  3 +++
 requirements.txt      |  3 +++
 3 files changed, 19 insertions(+)

diff --git a/mixnet/v2/sim/main.py b/mixnet/v2/sim/main.py
index 30bd22c..e4e9e0d 100644
--- a/mixnet/v2/sim/main.py
+++ b/mixnet/v2/sim/main.py
@@ -1,5 +1,9 @@
 import argparse
 
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn
+
 from simulation import Simulation
 
 if __name__ == "__main__":
@@ -11,4 +15,13 @@ if __name__ == "__main__":
 
     sim = Simulation(args.num_nodes, args.num_mix_layers)
     sim.run(until=args.running_time)
+
+    df = pd.DataFrame(sim.p2p.message_sizes, columns=["message_size"])
+    print(df.describe())
+    plt.figure(figsize=(10, 6))
+    seaborn.boxplot(y=df["message_size"])
+    plt.title("Message size distribution")
+    plt.ylabel("Message Size (bytes)")
+    plt.show()
+
     print("Simulation complete!")
\ No newline at end of file
diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py
index 4f5244e..792e943 100644
--- a/mixnet/v2/sim/p2p.py
+++ b/mixnet/v2/sim/p2p.py
@@ -9,6 +9,7 @@ class P2p:
     def __init__(self, env: simpy.Environment):
         self.env = env
         self.nodes = []
+        self.message_sizes = []
 
     def add_node(self, nodes):
         self.nodes.extend(nodes)
@@ -16,6 +17,8 @@ class P2p:
     # TODO: This should accept only bytes, but SphinxPacket is also accepted until we implement the Sphinx serde
     def broadcast(self, msg: SphinxPacket | bytes):
         self.log("Broadcasting a msg: %d bytes" % len(msg))
+        self.message_sizes.append(len(msg))
+
         yield self.env.timeout(1)
         # TODO: gossipsub or something similar
         for node in self.nodes:
diff --git a/requirements.txt b/requirements.txt
index 4edc555..13b4439 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,3 +8,6 @@ scipy==1.11.4
 black==23.12.1
 sympy==1.12
 simpy==4.1.1
+pandas==2.2.2
+matplotlib==3.8.4
+seaborn==0.13.2