plot_topology.py

"""Render the mine road network to ``topology.png`` straight from the data.

This is an optional visualisation. It reads ``data/nodes.csv`` and
``data/edges.csv`` via the model's own loaders so the figure can never drift
from the simulated network. Capacity-constrained (single-lane) segments are
drawn thick and red so the structural bottlenecks are visible at a glance.

Usage:
    python3 plot_topology.py [--out topology.png]
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # headless / reproducible
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

import mine_sim

HERE = Path(__file__).resolve().parent

# Marker style and colour per node type.
NODE_STYLE = {
    "parking": ("s", "#888888", "Parking"),
    "junction": ("o", "#4477aa", "Junction"),
    "load_ore": ("s", "#228833", "Ore loader"),
    "crusher": ("^", "#cc3311", "Crusher (dump)"),
    "waste_dump": ("v", "#aa7744", "Waste dump"),
    "maintenance": ("D", "#aa3377", "Maintenance"),
}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=HERE / "data")
    parser.add_argument("--out", type=Path, default=HERE / "topology.png")
    args = parser.parse_args(argv)

    nodes = mine_sim.load_nodes(args.data_dir)
    edges = mine_sim.load_edges(args.data_dir)

    fig, ax = plt.subplots(figsize=(12, 9))

    # Edges: constrained (capacity < 999) thick red, others thin grey.
    drawn_constrained = drawn_open = False
    for e in edges.values():
        a, b = nodes[e.from_node], nodes[e.to_node]
        if e.constrained:
            ax.annotate(
                "", xy=(b.x, b.y), xytext=(a.x, a.y),
                arrowprops=dict(arrowstyle="-|>", color="#cc3311", lw=2.2,
                                alpha=0.9, shrinkA=12, shrinkB=12),
            )
            drawn_constrained = True
        else:
            ax.annotate(
                "", xy=(b.x, b.y), xytext=(a.x, a.y),
                arrowprops=dict(arrowstyle="-|>", color="#bbbbbb", lw=1.0,
                                alpha=0.7, shrinkA=12, shrinkB=12),
            )
            drawn_open = True

    # Nodes.
    seen_types: set[str] = set()
    for n in nodes.values():
        marker, colour, _ = NODE_STYLE.get(n.node_type, ("o", "#000000", n.node_type))
        ax.scatter(n.x, n.y, marker=marker, c=colour, s=240, edgecolors="black",
                   zorder=3, linewidths=0.8)
        ax.annotate(f"{n.node_id}", (n.x, n.y), xytext=(0, 12),
                    textcoords="offset points", ha="center", fontsize=8, zorder=4)
        seen_types.add(n.node_type)

    # Legend.
    handles = [
        Line2D([0], [0], marker=NODE_STYLE[t][0], color="w",
               markerfacecolor=NODE_STYLE[t][1], markeredgecolor="black",
               markersize=11, label=NODE_STYLE[t][2])
        for t in NODE_STYLE if t in seen_types
    ]
    if drawn_constrained:
        handles.append(Line2D([0], [0], color="#cc3311", lw=2.2,
                              label="Capacity-1 segment (single lane)"))
    if drawn_open:
        handles.append(Line2D([0], [0], color="#bbbbbb", lw=1.0, label="Open road"))
    ax.legend(handles=handles, loc="upper left", framealpha=0.9, fontsize=9)

    ax.set_title("Synthetic mine road network — single-lane segments highlighted")
    ax.set_xlabel("x (m)")
    ax.set_ylabel("y (m)")
    ax.set_aspect("equal", adjustable="datalim")
    ax.grid(True, alpha=0.2)
    fig.tight_layout()
    fig.savefig(args.out, dpi=130)
    print(f"Wrote {args.out}")
    return 0


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