import folium
import geopandas as gpd
import numpy as np
import osmium
import pandas as pd
from shapely.geometry import Point, LineString, Polygon

np.set_printoptions(suppress=True, precision=8, linewidth=1000)
pd.options.mode.chained_assignment = None
pd.set_option("display.max_columns", None)
pd.set_option("display.width", None)


class FireStationHandler(osmium.SimpleHandler):
    """
    Collects OSM elements with amenity=fire_station.
    - Nodes: Point
    - Ways: Polygon if closed (>=4 nodes and first==last), else LineString
    """
    def __init__(self):
        super().__init__()
        self.features = []

    def node(self, n):
        if n.tags.get("amenity") == "fire_station" and n.location and n.location.valid():
            self.features.append({
                "osm_id": n.id,
                "osm_type": "node",
                "name": n.tags.get("name"),
                "lon": float(n.location.lon),
                "lat": float(n.location.lat),
                "tags": {t.k: t.v for t in n.tags},
            })

    def way(self, w):
        if w.tags.get("amenity") != "fire_station":
            return

        coords = []
        for nd in w.nodes:
            if not nd.location or not nd.location.valid():
                coords = []
                break
            coords.append((float(nd.location.lon), float(nd.location.lat)))

        if not coords:
            return

        lon = sum(c[0] for c in coords) / len(coords)
        lat = sum(c[1] for c in coords) / len(coords)
        
        self.features.append({
            "osm_id": w.id,
            "osm_type": "way",
            "name": w.tags.get("name"),
            "lon": lon,
            "lat": lat,
            "tags": {t.k: t.v for t in w.tags},
        })

# Path to the OSM PBF file.
pbf_path = "iowa-latest.osm.pbf"

# Apply the handler to the PBF file.
handler = FireStationHandler()
handler.apply_file(pbf_path, locations=True)  

# Create DataFrame from the extracted data.
df = pd.DataFrame(handler.features)

# Convert DataFrame to a GeoDataFrame.
df = gpd.GeoDataFrame(
    df, 
    geometry=gpd.points_from_xy(df.lon, df.lat), 
    crs="EPSG:4326"
)

print(f"\nExtracted {len(df)} fire stations from {pbf_path}.")

df.head()

Extracted 626 fire stations from iowa-latest.osm.pbf.

def is_fire_station(tags):
    """
    Determines if the given OSM tags indicate a fire station.
    """
    g = lambda k: (tags.get(k) or "").lower()
    return (
        g("amenity") == "fire_station"
        or g("office") == "fire_service"
        or g("emergency") in {"fire_service", "rescue_station"}
        or g("building") == "fire_station"
    )


class FireStationHandlerV2(osmium.SimpleHandler):
    """
    Collects OSM elements that are fire stations based on multiple tags.
    """
    def __init__(self):
        super().__init__()
        self.features = []

    def node(self, n):
        """
        Processes OSM nodes to identify fire stations.
        """
        if is_fire_station(n.tags) and n.location and n.location.valid():
            self.features.append({
                "osm_id": n.id,
                "osm_type": "node",
                "name": n.tags.get("name"),
                "lon": float(n.location.lon),
                "lat": float(n.location.lat),
                "tags": {t.k: t.v for t in n.tags},
            })

    def way(self, w):
        """
        Processes OSM ways to identify fire stations.
        """
        if not is_fire_station(w.tags): 
            return
        coords = []
        for nd in w.nodes:
            if not nd.location or not nd.location.valid():
                coords = []
                break
            coords.append((float(nd.location.lon), float(nd.location.lat)))
        if not coords: 
            return
        
        lon = sum(c[0] for c in coords) / len(coords)
        lat = sum(c[1] for c in coords) / len(coords)
        
        self.features.append({
            "osm_id": w.id,
            "osm_type": "way",
            "name": w.tags.get("name"),
            "lon": lon,
            "lat": lat,
            "tags": {t.k: t.v for t in w.tags},
        })

# Path to the OSM PBF file.
pbf_path = "iowa-latest.osm.pbf"

# Apply the handler to the PBF file.
handler = FireStationHandlerV2()
handler.apply_file(pbf_path, locations=True)  

# Create DataFrame from the extracted data.
df = pd.DataFrame(handler.features)

# Convert DataFrame to a GeoDataFrame.
df = gpd.GeoDataFrame(
    df, 
    geometry=gpd.points_from_xy(df.lon, df.lat), 
    crs="EPSG:4326"
)

print(f"\nExtracted {len(df)} fire stations from {pbf_path}.")

df.head()

Extracted 641 fire stations from iowa-latest.osm.pbf.

# Path to the OSM PBF file.
pbf_path = "iowa-firestations-filtered.osm.pbf"

# Apply the handler to the PBF file.
handler = FireStationHandlerV2()
handler.apply_file(pbf_path, locations=True)  

# Create DataFrame from the extracted data.
df = pd.DataFrame(handler.features)

# Convert DataFrame to a GeoDataFrame.
df = gpd.GeoDataFrame(
    df, 
    geometry=gpd.points_from_xy(df.lon, df.lat), 
    crs="EPSG:4326"
)

print(f"\nExtracted {len(df)} fire stations from {pbf_path}.")

df.head()

Extracted 626 fire stations from iowa-firestations-filtered.osm.pbf.

# Unpack the coordinates and names.
lats, lons, names = df.lat.tolist(), df.lon.tolist(), df.name.tolist()

# Center map on central Iowa. 
mid_lon, mid_lat = -93.0977, 41.8780
 
m = folium.Map(
    location=[mid_lat, mid_lon], 
    zoom_start=8
)

for name, lon, lat in zip(names, lons, lats):

    folium.CircleMarker(
        location=[lat, lon], 
        radius=3, 
        color="red", 
        fill_color="red", 
        popup=name,
        fill=True,
        fill_opacity=1
        ).add_to(m)

m