plot.py

import math
from typing import Optional
import tilemapbase as tmb
import numpy as np
from matplotlib import pyplot as plt, ticker
from matplotlib.collections import LineCollection
from numpy.typing import NDArray
from tdp import TravelingDeliverymanProblem
import numba as nb


def hover_annot(event, ax: plt.Axes, sc_nodes, annot,
                nodes_lat_lon, graph: TravelingDeliverymanProblem) -> None:
    """
    Update annotation on hover event to show useful information on geographic plot.
    """
    if event.inaxes == ax:
        cont, ind = sc_nodes.contains(event)
        if cont:
            idx = ind["ind"][0]
            pos = sc_nodes.get_offsets()[idx]
            annot.xy = pos
            deadline = graph.get_deadline_by_id(graph.node_id_by_idx(idx))
            if deadline is None:
                annot.set_text(f"Deadline: None\nLat: {nodes_lat_lon[:, 0][idx]:.4f}\nLon: {nodes_lat_lon[:, 1][idx]:.4f}")
            else:
                annot.set_text(f"Deadline: {deadline}\nLat: {nodes_lat_lon[:, 0][idx]:.4f}\nLon: {nodes_lat_lon[:, 1][idx]:.4f}")
            annot.set_visible(True)
            ax.figure.canvas.draw_idle()
        else:
            if annot.get_visible():
                annot.set_visible(False)
                ax.figure.canvas.draw_idle()


def fancy_geo_axis(ax: plt.Axes) -> None:
    """
    Format geographic axes with degree labels that start and stop at the start and stop of the graph area.
    """
    def format_x(x, pos):
        # We only care about longitude here.
        # We pass 0 for y because x determines longitude in Web Mercator independent of y
        lon, _ = tmb.to_lonlat(x, 0)
        return f"{lon:.4f}°"

    def format_y(y, pos):
        # We only care about latitude here.
        # We pass 0 for x because y determines latitude in Web Mercator independent of x
        _, lat = tmb.to_lonlat(0, y)
        return f"{lat:.4f}°"

    x0, x1 = ax.get_xlim()
    y0, y1 = ax.get_ylim()

    # 3. Create exactly 5 ticks per axis, evenly spaced from edge to edge
    # np.linspace ensures the first and last items are exactly the min/max
    num_ticks = 5
    xticks = np.linspace(x0, x1, num_ticks)
    yticks = np.linspace(y0, y1, num_ticks)

    # 4. Apply Locators (Positions) and Formatters (Labels)
    ax.xaxis.set_major_locator(ticker.FixedLocator(xticks))
    ax.yaxis.set_major_locator(ticker.FixedLocator(yticks))

    ax.xaxis.set_major_formatter(ticker.FuncFormatter(format_x))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(format_y))

    # 5. Make sure axes are visible and formatted
    ax.xaxis.set_visible(True)
    ax.yaxis.set_visible(True)
    ax.tick_params(axis='x', rotation=30)  # Rotate x-labels to avoid overlap

def deadline_to_color_discrete(deadline: float) -> str:
    """
    Map a deadline value to a discrete color from green (early) to red (late).
    """
    if deadline <= 5:
        return '#ff0000'    # Red
    elif deadline <= 20.1:
        return '#ff4500'    # '#ff6000'  # Red-Orange
    elif deadline <= 30.1:
        return '#ff8c00'    # '#FFb700'  # Orange
    elif deadline <= 40.1:
        return '#FFc800'    # Yellow
    elif deadline <= 60.1:
        return '#ccff00'    # Yellow-Green
    else:
        return '#00FF00'    # Green


@nb.njit("float64[:, :](float64[:], float64[:])", parallel=True, cache=True)
def lonlat_to_xy_mercator(lons: NDArray[np.float64], lats: NDArray[np.float64]) -> np.ndarray[np.float64]:
    """
    Project longitude/latitude coordinates (in degrees) into normalized
    Web Mercator (x, y) coordinates in the range [0, 1]. Optimized function for numba for fast native execution
    """
    n = lons.size
    out = np.empty((n, 2), dtype=np.float64)

    for i in nb.prange(n):
        lon = lons[i]
        lat = lats[i]

        x = (lon + 180.0) / 360.0
        lat_rad = math.radians(lat)
        y = (1.0 - math.log(math.tan(lat_rad) + 1.0 / math.cos(lat_rad)) / math.pi) / 2.0

        out[i, 0] = x
        out[i, 1] = y

    return out


def get_tmb_plotter(lon_min: float, lon_max: float, lat_min: float, lat_max: float, width_pixels: int = 800, margin: float = 0.0025) -> tmb.Plotter:
    """
    Create a TileMapBase plotter for the specified geographic extent.
    """
    extent = tmb.Extent.from_lonlat(
        lon_min - margin, lon_max + margin,
        lat_min - margin, lat_max + margin
    )
    # Get a plotter object (auto-calculates zoom based on desired pixel width)
    plotter = tmb.Plotter(extent, tmb.tiles.build_OSM(), width=width_pixels)
    return plotter

def geographic_plot(graph: TravelingDeliverymanProblem, ax: plt.Axes,
                    highlighted_tour: Optional[np.ndarray] = None,
                    tried_tour: Optional[np.ndarray] = None,
                    tour_color: str = 'red',
                    osm_background_alpha: float = 0.5,
                    sparse_view: bool = True) -> None:
    """
    Plot the geographic network using tilemapbase for geographic background suiting the coordinates.

    :param graph: TravelingDeliverymanProblem instance
    :param ax: Matplotlib Axes to plot on
    :param highlighted_tour: Optional tour to highlight (list of node indices)
    :param tried_tour: Optional tour to show as tried
    :param tour_color: Color for the highlighted tour
    :param osm_background_alpha: Opacity for the OSM background
    :param sparse_view: Limits the display to show only direct node connection and no synthetic edges with represent the real world connections
    """
    nodes_deadlined_lat_lon = [
        (data["lat"], data["lon"], data["deadline"])
        for _, data in graph.nodes(data=True)
        if "deadline" in data
    ]

    # [[lat, lon], [lat, lon], ...]
    nodes_lat_lon = np.asarray([graph.get_coordinates_by_id(node_id) for node_id in graph],
                               dtype=np.float64)  # shape (N, 2)
    # [[x, y], [x, y], ...]
    nodes_xy = lonlat_to_xy_mercator(nodes_lat_lon[:, 1], nodes_lat_lon[:, 0])  # shape (N, 2)

    # Process edges - handle synthetic edges with path_coords
    edges_xy_list = []
    for node1, node2, edge_data in graph.edges(data=True):
        if 'path_coords' in edge_data:
            if sparse_view:
                # path_coords is list of (y, x) = (lat, lon)
                path_coords = np.asarray(edge_data['path_coords'], dtype=np.float64)  # shape (P, 2)
                # Convert path to xy
                path_xy = lonlat_to_xy_mercator(path_coords[:, 1], path_coords[:, 0])  # shape (P, 2)
                # Create line segments from consecutive points
                for i in range(len(path_xy) - 1):
                    edges_xy_list.append([path_xy[i], path_xy[i + 1]])
            else:
                # Show direct connection not path, which is usually hidden by normal non synth edges
                coord1 = graph.get_coordinates_by_id(node1)
                coord2 = graph.get_coordinates_by_id(node2)
                xy1 = lonlat_to_xy_mercator(np.array([coord1[1]]), np.array([coord1[0]]))[0]
                xy2 = lonlat_to_xy_mercator(np.array([coord2[1]]), np.array([coord2[0]]))[0]
                edges_xy_list.append([xy1, xy2])
        else:
            # Regular edge - straight line
            coord1 = graph.get_coordinates_by_id(node1)
            coord2 = graph.get_coordinates_by_id(node2)
            xy1 = lonlat_to_xy_mercator(np.array([coord1[1]]), np.array([coord1[0]]))[0]
            xy2 = lonlat_to_xy_mercator(np.array([coord2[1]]), np.array([coord2[0]]))[0]
            edges_xy_list.append([xy1, xy2])

    edges_xy = np.asarray(edges_xy_list, dtype=np.float64) if edges_xy_list else np.empty((0, 2, 2))

    ax.clear()
    ax.set_aspect('equal')  # Geographic maps should have equal aspect ratio
    plotter = get_tmb_plotter(
        lon_min=nodes_lat_lon[:, 1].min(),
        lon_max=nodes_lat_lon[:, 1].max(),
        lat_min=nodes_lat_lon[:, 0].min(),
        lat_max=nodes_lat_lon[:, 0].max()
    )
    plotter.plot(ax=ax, alpha=osm_background_alpha)

    # Draw nodes
    sc_nodes = ax.scatter(nodes_xy[:, 0], nodes_xy[:, 1], c='steelblue', s=100, zorder=5, edgecolors='darkblue',
                          linewidths=1.0)
    if len(nodes_deadlined_lat_lon) > 0:
        for deadlined_node in nodes_deadlined_lat_lon:
            lat, lon, deadline = deadlined_node
            x, y = lonlat_to_xy_mercator(np.array([lon]), np.array([lat]))[0]
            color = deadline_to_color_discrete(deadline)
            ax.scatter(x, y, c=color, s=150, zorder=6, edgecolors='black', linewidths=1.0)
            ax.text(x, y, str(deadline), color='black', fontsize=9, fontweight='bold', ha='center', va='center',
                    zorder=7)

    # Draw edges
    lc_edges = LineCollection(edges_xy, colors='gray', linewidths=0.8, alpha=0.5, zorder=2)
    ax.add_collection(lc_edges)

    if highlighted_tour is not None:
        # Could be eased out, by assuming tour includes all nodes, but this is more general and allows for subtour visualization
        # [[lat, lon], [lat, lon], ...]
        tour_lat_lon = np.asarray([graph.get_coordinates_by_id(graph.node_id_by_idx(idx)) for idx in highlighted_tour],
                                  dtype=np.float64)  # shape (N_tour, 2)
        # [[x, y], [x, y], ...]
        tour_xy = lonlat_to_xy_mercator(tour_lat_lon[:, 1], tour_lat_lon[:, 0])  # shape (N_tour, 2)
        # [[edge1], [edge2], ...], with edge = [[x1, y1], [x2, y2]] => [[[x1, y1], [x2, y2]], ...], easier than above, since we already have x/y and the tour is in order
        tour_edges = np.asarray([(tour_xy[i], tour_xy[(i + 1) % len(tour_xy)]) for i in range(len(tour_xy))],
                                dtype=np.float64)  # shape (N_tour, 2, 2)

        # --- Add arrow pointing to the START node ---
        start_x, start_y = tour_xy[0]
        ax.annotate(
            "Start",
            xy=(start_x, start_y),
            xytext=(start_x, start_y - 0.000005),
            arrowprops=dict(arrowstyle="simple", color=tour_color, lw=0.02),
            fontsize=8,
            fontweight='bold',
            color=tour_color,
            zorder=10
        )

        # Calculate start points and direction vectors
        starts = tour_edges[:, 0, :]  # shape (N_tour, 2)
        ends = tour_edges[:, 1, :]  # shape (N_tour, 2)
        directions = ends - starts  # shape (N_tour, 2)

        # Plot edges as arrows using quiver
        ax.quiver(
            starts[:, 0], starts[:, 1],  # X, Y starting points
            directions[:, 0], directions[:, 1],  # U, V direction vectors
            angles='xy',
            scale_units='xy',
            scale=1,
            color=tour_color,
            width=0.006,  # Adjust arrow width
            alpha=0.8,
            zorder=4
        )

    # Plot tried tour if provided
    if tried_tour is not None:
        tried_tour_lat_lon = np.asarray([graph.get_coordinates_by_id(graph.node_id_by_idx(idx)) for idx in tried_tour],
                                        dtype=np.float64)  # shape (N_tour, 2)
        tried_tour_xy = lonlat_to_xy_mercator(tried_tour_lat_lon[:, 1], tried_tour_lat_lon[:, 0])  # shape (N_tour, 2)
        tried_tour_edges = np.asarray(
            [(tried_tour_xy[i], tried_tour_xy[(i + 1) % len(tried_tour_xy)]) for i in range(len(tried_tour_xy))],
            dtype=np.float64)  # shape (N_tour, 2, 2)

        lc_tried_tour = LineCollection(tried_tour_edges, colors='orange', linewidths=1.0, alpha=0.6, zorder=3)
        ax.add_collection(lc_tried_tour)

    annot = ax.annotate("", xy=(0, 0), xytext=(15, 15), textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w", alpha=0.9), arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)
    ax.figure.canvas.mpl_connect("motion_notify_event",
                                 hover_annot_wrapper := lambda event: hover_annot(event, ax, sc_nodes, annot,
                                                                                  nodes_lat_lon, graph))
    fancy_geo_axis(ax)