Working with GTFS

import sys
import zipfile
from pathlib import Path

import folium
import geopandas as gpd
import pandas as pd
import requests

# ruff: noqa: T201

sys.path.append(str(Path("..").resolve()))
from src.geoai.cookbook_functions import get_base_map

CITY_NAME = "Vienna"
BBOX = "16.335005,48.187854,16.400923,48.209995"
CACHE_DIR = Path("./gtfs_data")
OUTPUT_DIR = "./output"

# GitLab data project (cookbooks/data/geospatial-data), same LFS pattern as intersections.csv
GTFS_ZIP_URL = (
    "https://gitlab.tuwien.ac.at/api/v4/projects/15378/repository/files/"
    "vienna_gtfs.zip/raw?ref=main&lfs=true"
)

ROUTE_TYPE_LABELS = {
    0: "Tram",
    1: "Subway",
    2: "Rail",
    3: "Bus",
}

min_lon, min_lat, max_lon, max_lat = map(float, BBOX.split(","))
bounds = [[min_lat, min_lon], [max_lat, max_lon]]
center_lat = (min_lat + max_lat) / 2
center_lon = (min_lon + max_lon) / 2

CACHE_DIR.mkdir(parents=True, exist_ok=True)
zip_path = CACHE_DIR / "vienna_gtfs.zip"
feed_dir = CACHE_DIR / CITY_NAME

if not zip_path.exists():
    print(f"Downloading {GTFS_ZIP_URL}")
    response = requests.get(GTFS_ZIP_URL, timeout=120)
    response.raise_for_status()
    zip_path.write_bytes(response.content)
    print(f"Saved {zip_path}")
else:
    print(f"Using cached {zip_path}")

if not (feed_dir / "stops.txt").exists():
    feed_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path) as archive:
        archive.extractall(feed_dir)
    print(f"Extracted feed to {feed_dir}")

print("Files in the feed:")
for path in sorted(feed_dir.glob("*.txt")):
    print(f"  {path.name:20s}  {path.stat().st_size:8d} bytes")

gtfs = {}
for path in sorted(feed_dir.glob("*.txt")):
    name = path.stem
    gtfs[name] = pd.read_csv(path, dtype=str)

for name, table in gtfs.items():
    print(f"{name:16s}  rows={len(table):7d}  columns={list(table.columns)}")

display(gtfs["stops"].head())

routes = gtfs["routes"].copy()
routes["mode"] = routes["route_type"].astype(int).map(ROUTE_TYPE_LABELS).fillna("Other")
display(routes.head())
print("Routes by mode:")
print(routes["mode"].value_counts())

stops = gtfs["stops"].copy()
stops["stop_lat"] = stops["stop_lat"].astype(float)
stops["stop_lon"] = stops["stop_lon"].astype(float)
stops_gdf = gpd.GeoDataFrame(
    stops,
    geometry=gpd.points_from_xy(stops["stop_lon"], stops["stop_lat"]),
    crs="EPSG:4326",
)
print(f"Stops: {len(stops_gdf)} points, CRS={stops_gdf.crs}")

m_stops = get_base_map(center_lat, center_lon)
folium.Rectangle(
    bounds=bounds, color="#3388ff", fill=False, weight=2, popup="Cookbook bounding box"
).add_to(m_stops)

for _, row in stops_gdf.iterrows():
    folium.CircleMarker(
        location=[row.stop_lat, row.stop_lon],
        radius=3,
        color="#d62728",
        fill=True,
        fill_opacity=0.7,
        popup=folium.Popup(
            f"<b>{row.stop_name}</b><br>id={row.stop_id}", max_width=250
        ),
    ).add_to(m_stops)

m_stops

st = gtfs["stop_times"][["trip_id", "stop_id"]].drop_duplicates()
trip_route = gtfs["trips"][["trip_id", "route_id"]].drop_duplicates()
linked = st.merge(trip_route, on="trip_id").merge(
    routes[["route_id", "route_short_name", "mode"]],
    on="route_id",
)

per_stop = (
    linked.groupby("stop_id")
    .agg(
        n_routes=("route_id", "nunique"),
        route_names=(
            "route_short_name",
            lambda s: ", ".join(sorted({str(x) for x in s if pd.notna(x)})),
        ),
    )
    .reset_index()
    .merge(gtfs["stops"][["stop_id", "stop_name"]], on="stop_id")
    .sort_values("n_routes", ascending=False)
)

per_stop.head(10)

if "shapes" in gtfs:
    shapes = gtfs["shapes"].copy()
    shapes["shape_pt_lat"] = shapes["shape_pt_lat"].astype(float)
    shapes["shape_pt_lon"] = shapes["shape_pt_lon"].astype(float)
    shapes["shape_pt_sequence"] = shapes["shape_pt_sequence"].astype(int)
    shapes = shapes.sort_values(["shape_id", "shape_pt_sequence"])

    from shapely.geometry import LineString

    lines = []
    for shape_id, group in shapes.groupby("shape_id"):
        coords = list(zip(group["shape_pt_lon"], group["shape_pt_lat"], strict=True))
        if len(coords) >= 2:
            lines.append({"shape_id": shape_id, "geometry": LineString(coords)})

    shapes_gdf = gpd.GeoDataFrame(lines, crs="EPSG:4326")
    trip_shape = (
        gtfs["trips"][["route_id", "shape_id"]]
        .dropna(subset=["shape_id"])
        .drop_duplicates(subset=["shape_id"])
    )
    shapes_gdf = shapes_gdf.merge(trip_shape, on="shape_id", how="left").merge(
        routes[["route_id", "route_short_name"]], on="route_id", how="left"
    )

    sample = shapes_gdf.head(12)
    m_shapes = get_base_map(center_lat, center_lon)
    folium.Rectangle(bounds=bounds, color="#3388ff", fill=False, weight=2).add_to(
        m_shapes
    )
    for _, row in sample.iterrows():
        folium.GeoJson(
            row.geometry.__geo_interface__,
            style_function=lambda _x: {"color": "#ff7f00", "weight": 3},
            tooltip=f"route {row.get('route_short_name', row.shape_id)}",
        ).add_to(m_shapes)
    print(f"Plotted {len(sample)} of {len(shapes_gdf)} shapes")
    display(m_shapes)
else:
    print("No shapes.txt in this feed sample.")