ping_plot.py

#!/usr/bin/env python3
import argparse
import json
import os
import statistics
import sys
from dataclasses import dataclass
from typing import Dict, List, Tuple


@dataclass
class Stats:
    median: float
    p99: float
    max: float


def percentile(values: List[float], pct: float) -> float:
    if not values:
        raise ValueError("percentile requires non-empty values")
    if pct <= 0:
        return min(values)
    if pct >= 100:
        return max(values)
    ordered = sorted(values)
    k = (len(ordered) - 1) * (pct / 100.0)
    f = int(k)
    c = min(f + 1, len(ordered) - 1)
    if f == c:
        return ordered[f]
    return ordered[f] + (ordered[c] - ordered[f]) * (k - f)


def compute_stats(values: List[float]) -> Stats:
    return Stats(
        median=statistics.median(values),
        p99=percentile(values, 99),
        max=max(values),
    )


def sanitize_filename(name: str) -> str:
    safe = []
    for ch in name:
        if ch.isalnum() or ch in ("-", "_"):
            safe.append(ch)
        elif ch.isspace():
            safe.append("_")
    return "".join(safe).strip("_") or "plot"


def load_samples(input_path: str) -> Tuple[List[Dict[str, str]], Dict[str, List[Tuple[float, float]]]]:
    with open(input_path, "r", encoding="utf-8") as handle:
        payload = json.load(handle)
    locations = payload.get("locations")
    samples = payload.get("samples")
    if not isinstance(locations, list) or not isinstance(samples, list):
        raise RuntimeError("Input JSON must include locations and samples arrays")
    latencies: Dict[str, List[Tuple[float, float]]] = {
        loc["address"]: [] for loc in locations if isinstance(loc, dict)
    }
    for sample in samples:
        if not isinstance(sample, dict):
            continue
        address = sample.get("address")
        elapsed = sample.get("elapsed_s")
        latency = sample.get("latency_ms")
        if isinstance(address, str) and isinstance(elapsed, (int, float)) and isinstance(latency, (int, float)):
            latencies.setdefault(address, []).append((float(elapsed), float(latency)))
    return locations, latencies


def plot_overview(
    locations: List[Dict[str, str]],
    latencies: Dict[str, List[Tuple[float, float]]],
    output_path: str,
) -> None:
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise RuntimeError("matplotlib is required to plot results") from exc

    addresses = [loc.get("address", "") for loc in locations]
    labels = [loc.get("name", loc.get("address", "")) for loc in locations]
    fig, ax = plt.subplots(figsize=(10, 6))
    x_positions = list(range(len(addresses)))

    for i, addr in enumerate(addresses):
        samples = latencies.get(addr, [])
        if not samples:
            continue
        values = [value for _, value in samples]
        for j, latency in enumerate(values):
            jitter = ((j % 5) - 2) * 0.05
            ax.scatter(i + jitter, latency, marker="x", color="tab:blue", alpha=0.7)
        stats = compute_stats(values)
        ax.scatter(i, stats.median, marker="o", color="tab:green", label="median" if i == 0 else "")
        ax.scatter(i, stats.p99, marker="^", color="tab:orange", label="p99" if i == 0 else "")
        ax.scatter(i, stats.max, marker="s", color="tab:red", label="max" if i == 0 else "")
        ax.annotate(
            f"{stats.median:.2f}",
            (i, stats.median),
            textcoords="offset points",
            xytext=(6, -6),
            fontsize=8,
            color="tab:green",
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white", edgecolor="none", alpha=0.8),
        )
        ax.annotate(
            f"{stats.p99:.2f}",
            (i, stats.p99),
            textcoords="offset points",
            xytext=(6, -6),
            fontsize=8,
            color="tab:orange",
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white", edgecolor="none", alpha=0.8),
        )

    ax.set_xticks(x_positions)
    ax.set_xticklabels(labels, rotation=30, ha="right")
    ax.set_ylabel("Latency (ms)")
    ax.set_xlabel("Location")
    ax.set_title("Ping Latency by Location")
    ax.set_ylim(0, 30)
    ax.grid(axis="y", linestyle="--", alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)


def plot_time_series(
    locations: List[Dict[str, str]],
    latencies: Dict[str, List[Tuple[float, float]]],
    output_dir: str,
) -> None:
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise RuntimeError("matplotlib is required to plot results") from exc

    for loc in locations:
        if not isinstance(loc, dict):
            continue
        label = loc.get("name", "")
        address = loc.get("address", "")
        samples = latencies.get(address, [])
        if not samples:
            continue
        times = [t for t, _ in samples]
        values = [v for _, v in samples]
        fig, ax = plt.subplots(figsize=(10, 4))
        ax.plot(times, values, marker="x", linestyle="-", color="tab:blue")
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Latency (ms)")
        ax.set_title(f"Latency Over Time: {label}")
        ax.set_ylim(0, 30)
        ax.grid(axis="y", linestyle="--", alpha=0.3)
        fig.tight_layout()
        filename = f"{sanitize_filename(label)}.png"
        fig.savefig(os.path.join(output_dir, filename), dpi=150)
        plt.close(fig)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Plot latency samples from ping data.")
    parser.add_argument(
        "--input-json",
        default="ping_output/ping_samples.json",
        help="Path to the JSON data file produced by ping_collect.py",
    )
    parser.add_argument(
        "--output-dir",
        default="ping_output",
        help="Output directory for plots",
    )
    return parser


def parse_args() -> argparse.Namespace:
    return build_parser().parse_args()


def main() -> int:
    if len(sys.argv) == 1:
        build_parser().print_help()
        return 2
    args = parse_args()
    locations, latencies = load_samples(args.input_json)
    os.makedirs(args.output_dir, exist_ok=True)
    overview_path = os.path.join(args.output_dir, "latency_overview.png")
    plot_overview(locations, latencies, overview_path)
    plot_time_series(locations, latencies, args.output_dir)
    print(f"Saved plots to {args.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())