Allow specifying the output plot filename via CLI

plot_propagation.py

@@ -9,6 +9,7 @@
 import sys
 import os
 import re
+import argparse
 from datetime import datetime
 from typing import List, Tuple
 import matplotlib.pyplot as plt
@@ -63,7 +64,7 @@ def parse_shadow_logs(node_count: int) -> List[Tuple[float, int]]:
     return cumulative_data
 
 
-def plot_propagation(node_count: int):
+def plot_propagation(node_count: int, output_file: str = 'message_propagation.png'):
     """Plot message propagation over time."""
     print(f"Parsing Shadow logs for {node_count} nodes...")
 
@@ -96,7 +97,6 @@ def plot_propagation(node_count: int):
     plt.tight_layout()
 
     # Save plot
-    output_file = 'message_propagation.png'
     plt.savefig(output_file, dpi=150)
     print(f"\nPlot saved to: {output_file}")
 
@@ -117,26 +117,33 @@ def plot_propagation(node_count: int):
 
 
 def main():
-    if len(sys.argv) != 2:
-        print("Usage: python3 plot_propagation.py <node-count>")
-        print("  node-count: Number of nodes in the simulation")
-        sys.exit(1)
-
-    try:
-        node_count = int(sys.argv[1])
-    except ValueError:
-        print("Error: node-count must be an integer")
-        sys.exit(1)
-
-    if node_count <= 0:
-        print("Error: node-count must be positive")
-        sys.exit(1)
+    parser = argparse.ArgumentParser(
+        description='Plot message propagation over time from Shadow simulation logs',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python3 plot_propagation.py 10
+  python3 plot_propagation.py 20 -o propagation_20nodes.png
+  python3 plot_propagation.py 10 --output results/test1.png
+        """
+    )
+
+    parser.add_argument('node_count', type=int, 
+                        help='Number of nodes in the simulation')
+    parser.add_argument('-o', '--output', type=str, 
+                        default='message_propagation.png',
+                        help='Output file path for the plot (default: message_propagation.png)')
+
+    args = parser.parse_args()
+
+    if args.node_count <= 0:
+        parser.error("node-count must be positive")
 
     if not os.path.exists("shadow.data"):
         print("Error: shadow.data directory not found. Run simulation first with 'make run-sim'")
         sys.exit(1)
 
-    plot_propagation(node_count)
+    plot_propagation(args.node_count, args.output)
 
 
 if __name__ == "__main__":