netem-trace 0.4.4

#!/usr/bin/env python3
"""
Example Python script demonstrating how to plot trace series exported from Rust.

This script reads the JSON/CSV files generated by the plot_traces example
and creates visualizations using matplotlib.

Requirements:
    pip install matplotlib polars

Usage:
    python examples/plot_example.py
"""

import json
import polars as pl
import matplotlib.pyplot as plt


def plot_bw_trace_from_json(filename, title):
    """Plot bandwidth trace from JSON file."""
    with open(filename, "r") as f:
        data = json.load(f)

    # Extract data points
    times = []
    bandwidths = []

    for point in data:
        start_time = point["start_time"]
        duration = point["duration"]
        bw_mbps = point["value"]["bps"] / 1_000_000

        # Create step plot by adding points at start and end of each segment
        times.append(start_time)
        bandwidths.append(bw_mbps)
        times.append(start_time + duration)
        bandwidths.append(bw_mbps)

    plt.figure(figsize=(10, 5))
    plt.plot(times, bandwidths, linewidth=2)
    plt.xlabel("Time (seconds)")
    plt.ylabel("Bandwidth (Mbps)")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    return plt


def plot_bw_trace_from_csv(filename, title):
    """Plot bandwidth trace from CSV file."""
    df = pl.read_csv(filename)

    # Create step plot
    times = []
    bandwidths = []

    for row in df.iter_rows(named=True):
        start_time = row["start_time_secs"]
        duration = row["duration_secs"]
        bw_mbps = row["bandwidth_bps"] / 1_000_000

        times.append(start_time)
        bandwidths.append(bw_mbps)
        times.append(start_time + duration)
        bandwidths.append(bw_mbps)

    plt.figure(figsize=(10, 5))
    plt.plot(times, bandwidths, linewidth=2)
    plt.xlabel("Time (seconds)")
    plt.ylabel("Bandwidth (Mbps)")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    return plt


def plot_delay_per_packet_from_csv(filename, title):
    """Plot per-packet delay from CSV file."""
    df = pl.read_csv(filename)

    plt.figure(figsize=(10, 5))
    plt.plot(
        df["packet_index"].to_list(),
        (df["delay_secs"] * 1000).to_list(),
        "o-",
        linewidth=2,
    )
    plt.xlabel("Packet Index")
    plt.ylabel("Delay (ms)")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    return plt


def create_multi_plot():
    """Create a multi-panel plot showing different trace types."""
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))

    # Plot 1: Static bandwidth
    with open("static_bw.json", "r") as f:
        data = json.load(f)
    times, bandwidths = [], []
    for point in data:
        start = point["start_time"]
        dur = point["duration"]
        bw = point["value"]["bps"] / 1_000_000
        times.extend([start, start + dur])
        bandwidths.extend([bw, bw])
    axes[0, 0].plot(times, bandwidths, linewidth=2)
    axes[0, 0].set_xlabel("Time (s)")
    axes[0, 0].set_ylabel("Bandwidth (Mbps)")
    axes[0, 0].set_title("Static Bandwidth Trace")
    axes[0, 0].grid(True, alpha=0.3)

    # Plot 2: Sawtooth bandwidth
    df = pl.read_csv("sawtooth_bw.csv")
    times, bandwidths = [], []
    for row in df.iter_rows(named=True):
        start = row["start_time_secs"]
        dur = row["duration_secs"]
        bw = row["bandwidth_bps"] / 1_000_000
        times.extend([start, start + dur])
        bandwidths.extend([bw, bw])
    axes[0, 1].plot(times, bandwidths, linewidth=2, color="orange")
    axes[0, 1].set_xlabel("Time (s)")
    axes[0, 1].set_ylabel("Bandwidth (Mbps)")
    axes[0, 1].set_title("Sawtooth Bandwidth Trace")
    axes[0, 1].grid(True, alpha=0.3)

    # Plot 3: Normal bandwidth
    df = pl.read_csv("normal_bw.csv")
    times, bandwidths = [], []
    for row in df.iter_rows(named=True):
        start = row["start_time_secs"]
        dur = row["duration_secs"]
        bw = row["bandwidth_bps"] / 1_000_000
        times.extend([start, start + dur])
        bandwidths.extend([bw, bw])
    axes[1, 0].plot(times, bandwidths, linewidth=2, color="green")
    axes[1, 0].set_xlabel("Time (s)")
    axes[1, 0].set_ylabel("Bandwidth (Mbps)")
    axes[1, 0].set_title("Normalized Bandwidth Trace")
    axes[1, 0].grid(True, alpha=0.3)

    # Plot 4: Repeated pattern
    df = pl.read_csv("repeated_bw.csv")
    times, bandwidths = [], []
    for row in df.iter_rows(named=True):
        start = row["start_time_secs"]
        dur = row["duration_secs"]
        bw = row["bandwidth_bps"] / 1_000_000
        times.extend([start, start + dur])
        bandwidths.extend([bw, bw])
    axes[1, 1].plot(times, bandwidths, linewidth=2, color="red")
    axes[1, 1].set_xlabel("Time (s)")
    axes[1, 1].set_ylabel("Bandwidth (Mbps)")
    axes[1, 1].set_title("Repeated Pattern Trace")
    axes[1, 1].grid(True, alpha=0.3)

    plt.tight_layout()
    return plt


if __name__ == "__main__":
    print("Creating plots from exported trace data...")

    # Create individual plots
    print("1. Plotting static bandwidth trace from JSON...")
    plt1 = plot_bw_trace_from_json("static_bw.json", "Static Bandwidth Trace")
    plt1.savefig("plot_static_bw.png", dpi=150)
    print("   Saved to plot_static_bw.png")

    print("2. Plotting sawtooth bandwidth trace from CSV...")
    plt2 = plot_bw_trace_from_csv("sawtooth_bw.csv", "Sawtooth Bandwidth Trace")
    plt2.savefig("plot_sawtooth_bw.png", dpi=150)
    print("   Saved to plot_sawtooth_bw.png")

    print("3. Plotting per-packet delay from CSV...")
    plt3 = plot_delay_per_packet_from_csv("per_packet_delay.csv", "Per-Packet Delay")
    plt3.savefig("plot_per_packet_delay.png", dpi=150)
    print("   Saved to plot_per_packet_delay.png")

    print("4. Creating multi-panel plot...")
    plt4 = create_multi_plot()
    plt4.savefig("plot_multi.png", dpi=150)
    print("   Saved to plot_multi.png")

    print("\nAll plots created successfully!")
    print("You can view the plots by opening the PNG files in your image viewer.")