jetlag-maps/copenhagen/scripts/prepare.py

#!/usr/bin/env python3
"""Merge GTFS-derived layers into a single master GeoPackage, enriched for rendering.

Reads (from data/processed):
    gtfs_shapes.gpkg   route-tagged line geometries (gtfs_to_geopackage.py)
    gtfs_stops.gpkg    raw stop poles (one row per physical pole)
    route_colors.csv   route metadata (agency_id, short/long name, type, colour)
    area.gpkg          City Pass area polygon (build_area.py, incl. waterways)
and from data/raw/gtfs: agency.txt, trips.txt, stop_times.txt.

Writes data/processed/master.gpkg with layers:
    lines     one feature per (style, ref, direction) — the shape that serves
              the most in-area stops (so drawn lines pass the drawn stops)
    stops     bus/ferry stops, one point per stop name, pruned to stops
              actually served by a drawn line (see STOP_LINE_MARGIN_M)
    stations  rail-family stations, one point per (name, style)

Stop names in the Rejseplanen feed are unambiguous per location (verified:
max spread within a name is ~350 m), so every merge is key-based on names —
no distance-based clustering anywhere.

Style classification comes from config/modes.yaml (agency + route_type);
file order is the priority when a pole is served by several modes.
Colours: GTFS route_color -> styling.yaml palette (the feed leaves
route_color empty for the Copenhagen operators, so the palette wins).
"""
import re
import sys

import geopandas as gpd
import pandas as pd
import yaml

from _common import CONFIG, GTFS_RAW, PROCESSED

LENGTH_CRS = "EPSG:25832"  # UTM 32N — metres, for length/centroid computations
STATION_STYLES = {"metro", "s_tog", "light_rail", "regional"}
# A stop is drawn only if a drawn line of one of its serving refs passes
# within this distance. Routes have variants; without this check a stop can
# end up further from the map than the variant we chose not to draw.
STOP_LINE_MARGIN_M = 300


def load_modes():
    return yaml.safe_load((CONFIG / "modes.yaml").read_text())["modes"]


def load_styling():
    return yaml.safe_load((CONFIG / "styling.yaml").read_text())


def load_area():
    area = gpd.read_file(PROCESSED / "area.gpkg")
    return area.geometry.union_all()


def classify(agency, route_type, modes):
    """Map a GTFS (agency_name, route_type) pair to a style key, or None."""
    for style, cfg in modes.items():
        if agency not in cfg.get("agencies", []):
            continue
        types = cfg.get("route_types")
        if types and route_type not in types:
            continue
        return style
    return None


def load_routes(modes):
    """route_colors.csv joined with agency names and classified by style."""
    agencies = pd.read_csv(GTFS_RAW / "agency.txt", dtype=str)
    agency_names = dict(zip(agencies["agency_id"], agencies["agency_name"]))
    routes = pd.read_csv(PROCESSED / "route_colors.csv", dtype=str)
    routes["route_type"] = pd.to_numeric(routes["route_type"], errors="coerce")
    routes["agency_name"] = routes["agency_id"].map(agency_names)
    routes["style"] = [
        classify(a, t, modes)
        for a, t in zip(routes["agency_name"], routes["route_type"])
    ]
    return routes


def categorise_bus(ref, patterns):
    s = "" if ref is None else str(ref)
    for cat, pat in patterns.items():
        if re.search(pat, s):
            return cat
    return "regular"


def resolve_colour(style, ref, bus_category, gtfs_colour, palette):
    # 1. GTFS route_color (mostly empty in this feed)
    if isinstance(gtfs_colour, str) and gtfs_colour.strip():
        return gtfs_colour
    # 2. palette
    p = palette.get(style, {})
    if style == "bus":
        return p.get(bus_category) or p.get("default")
    return p.get(ref) or p.get("default")


def build_lines(routes, styling, area_geom, trips, st, area_stop_ids):
    """One feature per (style, ref, direction).

    Geometry per group: the shape that serves the most in-area stops
    (tie-break: longest). Picking by length alone can draw a variant that
    skips stops shown on the map (terminal stubs, short-turn branches).
    """
    shapes = gpd.read_file(PROCESSED / "gtfs_shapes.gpkg")

    shape_rows = trips.dropna(subset=["shape_id"]).drop_duplicates("shape_id")
    shape_direction = dict(zip(shape_rows["shape_id"], shape_rows["direction_id"]))

    # coverage: distinct in-area stops visited per shape
    trip_shape = dict(zip(trips["trip_id"], trips["shape_id"]))
    sv = st.assign(shape_id=st["trip_id"].map(trip_shape))
    sv = sv[sv["stop_id"].isin(area_stop_ids)]
    coverage = sv.groupby("shape_id")["stop_id"].nunique()

    use = routes.dropna(subset=["style"])[
        ["route_id", "route_short_name", "route_long_name",
         "route_color", "style"]
    ]
    g = shapes.merge(use, on="route_id", how="inner")
    g["direction_id"] = g["shape_id"].map(shape_direction)

    g = g.to_crs(LENGTH_CRS)
    g["_len"] = g.geometry.length
    g["_cov"] = g["shape_id"].map(coverage).fillna(0)
    g = (
        g.sort_values(["_cov", "_len"], ascending=False)
        .drop_duplicates(["style", "route_short_name", "direction_id"])
        .to_crs("EPSG:4326")
    )

    g["ref"] = g["route_short_name"]
    g["name"] = [
        ln if isinstance(ln, str) and ln.strip() else ref
        for ln, ref in zip(g["route_long_name"], g["ref"])
    ]
    patterns = styling["bus_filters"]["categories"]
    g["bus_category"] = [
        categorise_bus(ref, patterns) if style == "bus" else None
        for ref, style in zip(g["ref"], g["style"])
    ]
    palette = styling["palette"]
    g["colour_final"] = [
        resolve_colour(style, ref, cat, gtfs_c, palette)
        for style, ref, cat, gtfs_c
        in zip(g["style"], g["ref"], g["bus_category"], g["route_color"])
    ]

    g = g[["ref", "name", "style", "bus_category", "colour_final", "geometry"]]
    g = g[~g.geometry.isna() & g.geometry.is_valid]
    # the area polygon includes patched-in waterways (area.json), so ferry
    # shapes and sub-harbour metro tunnels survive the clip unfragmented
    return gpd.clip(g, area_geom)


def build_pole_classes(routes, modes, trips, st, area_stop_ids):
    """Per-pole classification and serving refs, from stop_times.

    Returns (pole_styles, pole_refs):
      pole_styles[stop_id] = highest-priority style among classified serving
                             routes (modes.yaml order)
      pole_refs[stop_id]   = set of serving route short names (classified)
    """
    classified = routes.dropna(subset=["style"])
    route_style = dict(zip(classified["route_id"], classified["style"]))
    route_ref = dict(
        zip(classified["route_id"], classified["route_short_name"])
    )
    trip_route = dict(zip(trips["trip_id"], trips["route_id"]))

    sv = st[st["stop_id"].isin(area_stop_ids)].copy()
    sv["route_id"] = sv["trip_id"].map(trip_route)
    sv = sv.dropna(subset=["route_id"])
    sv["style"] = sv["route_id"].map(route_style)
    sv = sv.dropna(subset=["style"])

    order = {s: i for i, s in enumerate(modes)}
    sv["_p"] = sv["style"].map(order)
    best = sv.loc[sv.groupby("stop_id")["_p"].idxmin(), ["stop_id", "style"]]
    pole_styles = dict(zip(best["stop_id"], best["style"]))

    pole_refs = sv.groupby("stop_id")["route_id"].apply(
        lambda ids: {route_ref[i] for i in ids}
    ).to_dict()
    return pole_styles, pole_refs


def build_stops_and_stations(pole_styles, pole_refs, area_geom, lines):
    poles = gpd.read_file(PROCESSED / "gtfs_stops.gpkg")
    poles["style"] = poles["stop_id"].map(pole_styles)
    poles["refs"] = poles["stop_id"].map(lambda s: pole_refs.get(s, set()))
    poles = poles.dropna(subset=["style"])

    # project for accurate centroids and distances
    poles = poles.to_crs(LENGTH_CRS)

    # stops (everything not rail-family: buses + harbour ferries), one row
    # per name, union of serving refs across its poles
    not_rail = poles[~poles["style"].isin(STATION_STYLES)]
    stop_rows = []
    for (name, style), grp in not_rail.groupby(["stop_name", "style"]):
        refs = set().union(*grp["refs"]) if len(grp) else set()
        stop_rows.append({
            "name": name,
            "style": style,
            "n_poles": len(grp),
            "refs": refs,
            "geometry": grp.geometry.union_all().centroid,
        })

    # stations: one row per (name, style)
    rail = poles[poles["style"].isin(STATION_STYLES)]
    station_rows = []
    for (name, style), grp in rail.groupby(["stop_name", "style"]):
        station_rows.append({
            "name": name,
            "style": style,
            "n_poles": len(grp),
            "geometry": grp.geometry.union_all().centroid,
        })

    area_25832 = gpd.GeoSeries([area_geom], crs="EPSG:4326").to_crs(LENGTH_CRS).union_all()

    stops = gpd.GeoDataFrame(stop_rows, crs=LENGTH_CRS)

    # prune stops not served by any drawn line within STOP_LINE_MARGIN_M
    if len(stops) and len(lines):
        line_geom = (
            lines.to_crs(LENGTH_CRS)
            .groupby("ref")["geometry"]
            .agg(lambda g: g.union_all())
            .to_dict()
        )
        keep = []
        for _, r in stops.iterrows():
            dists = [
                geom.distance(r.geometry)
                for rf in r["refs"]
                if (geom := line_geom.get(rf)) is not None
            ]
            keep.append(bool(dists) and min(dists) <= STOP_LINE_MARGIN_M)
        n_dropped = (~pd.Series(keep, index=stops.index)).sum()
        if n_dropped:
            print(f"pruned {n_dropped} stops not served by a drawn line "
                  f"(>{STOP_LINE_MARGIN_M} m from nearest)", flush=True)
        stops = stops[keep].drop(columns=["refs"])
    else:
        stops = stops.drop(columns=["refs"])

    stops = gpd.clip(stops, area_25832)
    stations = gpd.clip(
        gpd.GeoDataFrame(station_rows, crs=LENGTH_CRS), area_25832
    )
    return stops.to_crs("EPSG:4326"), stations.to_crs("EPSG:4326")


def main():
    for dep in ("gtfs_shapes.gpkg", "gtfs_stops.gpkg", "route_colors.csv"):
        if not (PROCESSED / dep).exists():
            print(f"prepare: missing {PROCESSED / dep} — "
                  "run gtfs_to_geopackage.py first.", file=sys.stderr)
            sys.exit(2)

    modes = load_modes()
    styling = load_styling()
    area_geom = load_area()

    routes = load_routes(modes)
    print(f"routes classified: {routes['style'].notna().sum()} of "
          f"{len(routes)} map to a style", flush=True)

    area_stop_ids = set(
        gpd.read_file(PROCESSED / "gtfs_stops.gpkg")["stop_id"]
    )
    trips = pd.read_csv(
        GTFS_RAW / "trips.txt",
        usecols=["trip_id", "route_id", "shape_id", "direction_id"],
        dtype=str,
    )
    st = pd.read_csv(
        GTFS_RAW / "stop_times.txt",
        usecols=["trip_id", "stop_id"],
        dtype=str,
    )

    lines = build_lines(routes, styling, area_geom, trips, st, area_stop_ids)
    print(f"lines: {len(lines)} features (one per style/ref/direction)",
          flush=True)

    pole_styles, pole_refs = build_pole_classes(
        routes, modes, trips, st, area_stop_ids
    )
    stops, stations = build_stops_and_stations(
        pole_styles, pole_refs, area_geom, lines
    )
    print(f"stops: {len(stops)} (bus + ferry); stations: {len(stations)}",
          flush=True)

    out = PROCESSED / "master.gpkg"
    if out.exists():
        out.unlink()
    lines.to_file(out, driver="GPKG", layer="lines")
    print(f"wrote lines layer: {len(lines)} features", flush=True)
    stops.to_file(out, driver="GPKG", layer="stops")
    print(f"wrote stops layer: {len(stops)} features", flush=True)
    stations.to_file(out, driver="GPKG", layer="stations")
    print(f"wrote stations layer: {len(stations)} features", flush=True)

    print("\nsummary by style:")
    for style, grp in lines.groupby("style"):
        refs = grp["ref"].dropna().unique()
        print(f"  {style:10s}: {len(grp):4d} feats, {len(refs):3d} refs")


if __name__ == "__main__":
    main()