OpenStreetMap extracts¶

from osmfinder import OsmExtractSource, display_available_extracts

from quackosm import convert_osm_extract_to_geodataframe

from shapely import from_wkt
from shapely.geometry import shape

from osmfinder import find_extracts_by_geometry
from quackosm import convert_geometry_to_geodataframe, geocode_to_geometry

import contextily as cx
import geopandas as gpd
from matplotlib import patches as mpatches
from matplotlib import pyplot as plt
from shapely import Polygon


def plot_full_geometry_coverage_breakdown(result) -> None:
    """Visualize the step-by-step geometry coverage breakdown.

    Steps are grouped by sub_polygon_index. Within each group, discarded
    steps are reduced to a single representative (the first one encountered).
    Each group is plotted in its own figure.
    """
    steps = result.steps
    if len(steps) == 0:
        return

    # Group steps by sub_polygon_index
    from collections import OrderedDict
    groups: OrderedDict[int | None, list] = OrderedDict()
    for step in steps:
        key = step.sub_polygon_index
        if key not in groups:
            groups[key] = []
        groups[key].append(step)

    for poly_idx, raw_group_steps in groups.items():
        group_steps = raw_group_steps
        # Reduce discarded steps within this group: keep only one
        total_group_steps = len(group_steps)
        discarded_indices = [i for i, s in enumerate(group_steps) if not s.selected]
        discarded_to_remove = max(0, len(discarded_indices) - 1)
        if discarded_to_remove > 0:
            keep_discarded_idx = discarded_indices[0]
            remove_set = set(discarded_indices[1:])
            group_steps = [s for i, s in enumerate(group_steps) if i not in remove_set]
            total_group_steps = len(group_steps)
            discarded_label_idx = keep_discarded_idx
        else:
            discarded_label_idx = -1

        if total_group_steps == 0:
            continue

        fig, axes = plt.subplots(nrows=total_group_steps, ncols=4, figsize=(20, 5 * total_group_steps))

        geometry_to_cover = group_steps[0].geometry_to_cover
        close_up_bbox = geometry_to_cover.bounds

        full_bbox = (
            gpd.GeoSeries([s.extract.geometry for s in group_steps], crs=4326)
            .union_all()
            .bounds
        )

        for ax_idx, step in enumerate(group_steps):
            iou_metric = step.iou
            is_selected = step.selected
            is_redundant = not is_selected and step.reason == "redundant"
            color = "C2" if is_selected else ("C8" if is_redundant else "C3")
            status = "accepted" if is_selected else ("redundant" if is_redundant else "discarded")

            if total_group_steps > 1:
                combined_ax = axes[ax_idx, 0]
                geometry_to_cover_ax = axes[ax_idx, 1]
                intersection_geometry_ax = axes[ax_idx, 2]
                extract_geometry_ax = axes[ax_idx, 3]
            else:
                combined_ax = axes[0]
                geometry_to_cover_ax = axes[1]
                intersection_geometry_ax = axes[2]
                extract_geometry_ax = axes[3]

            combined_ax.set_xlim([full_bbox[0], full_bbox[2]])
            combined_ax.set_ylim([full_bbox[1], full_bbox[3]])

            combined_ax.set_yticklabels([])
            combined_ax.set_xticklabels([])
            combined_ax.set_xticks([])
            combined_ax.set_yticks([])

            cx.add_basemap(combined_ax, source=cx.providers.Esri.WorldGrayCanvas, crs=4326)

            for ax in (geometry_to_cover_ax, intersection_geometry_ax, extract_geometry_ax):
                ax.set_xlim([close_up_bbox[0], close_up_bbox[2]])
                ax.set_ylim([close_up_bbox[1], close_up_bbox[3]])

                ax.set_axis_off()

                cx.add_basemap(ax, source=cx.providers.Esri.WorldGrayCanvas, crs=4326)

            combined_ax.set_title(step.extract.file_name)
            prefix = f"[polygon {poly_idx}] " if poly_idx is not None else ""
            if ax_idx == 0 and is_selected:
                combined_ax.set_ylabel(f"{prefix}Step {ax_idx + 1} (always accepted)")
            else:
                label = f"{prefix}Step {ax_idx + 1} ({status})"
                if ax_idx == discarded_label_idx and discarded_to_remove > 0:
                    label += f" (+ {discarded_to_remove} discarded)"
                combined_ax.set_ylabel(label)

            geometry_to_cover_ax.set_title("Geometry to cover")
            extract_geometry_ax.set_title("Extract geometry close-up")
            intersection_geometry_ax.set_title(f"Intersection (IoU: {iou_metric:.3g})")

            extract_geometry = step.extract.geometry
            intersection_geometry = step.intersection_geometry
            gpd.GeoSeries([step.geometry_to_cover], crs=4326).plot(
                ax=geometry_to_cover_ax, color="C0", alpha=0.8
            )
            gpd.GeoSeries([step.geometry_to_cover], crs=4326).plot(ax=combined_ax, color="C0", alpha=0.8)

            gpd.GeoSeries([extract_geometry], crs=4326).plot(
                ax=combined_ax,
                color=color,
                alpha=0.4,
                zorder=1,
            )
            gpd.GeoSeries([extract_geometry], crs=4326).plot(
                ax=extract_geometry_ax,
                color=color,
                alpha=0.4,
                zorder=1,
            )

            gpd.GeoSeries([intersection_geometry], crs=4326).plot(
                ax=combined_ax,
                color=(0, 0, 0, 0),
                zorder=2,
                hatch="///",
                edgecolor="C1",
                linewidth=1.5,
            )

            gpd.GeoSeries([intersection_geometry], crs=4326).plot(
                ax=intersection_geometry_ax,
                color=(0, 0, 0, 0),
                zorder=2,
                hatch="///",
                edgecolor="C1",
                linewidth=1.5,
            )

        fig.suptitle(f"Polygon {poly_idx}" if poly_idx is not None else "Geometry coverage", fontsize=14)
        fig.tight_layout()
        plt.show()


def plot_features_with_geometry_filter(
    features_gdf: gpd.GeoDataFrame, geometry_filter: Polygon
) -> None:
    fig = plt.figure()
    ax = fig.subplots()
    close_up_bbox = geometry_filter.bounds
    ax.set_xlim([close_up_bbox[0], close_up_bbox[2]])
    ax.set_ylim([close_up_bbox[1], close_up_bbox[3]])

    ax.set_axis_off()

    cx.add_basemap(ax, source=cx.providers.Esri.WorldGrayCanvas, crs=4326)

    features_gdf.plot(ax=ax, markersize=1, zorder=1, alpha=0.2)
    features_gdf.boundary.plot(ax=ax, markersize=0, zorder=1, alpha=0.2)
    gpd.GeoSeries([geometry_filter], crs=4326).plot(
        ax=ax,
        color=(0, 0, 0, 0),
        zorder=2,
        hatch="///",
        edgecolor="orange",
        linewidth=1.5,
    )

    blue_patch = mpatches.Patch(color="C0", alpha=0.8, label="OSM features")
    orange_patch = mpatches.Patch(
        facecolor=(0, 0, 0, 0),
        edgecolor="orange",
        hatch="///",
        linewidth=1.5,
        label="Geometry filter",
    )
    ax.legend(handles=[blue_patch, orange_patch], loc="lower right")

    fig.tight_layout()

    plt.show()