src/data_processing/extract_osm_data.py

"""
Extract infrastructure data from OpenStreetMap via Overpass API.
Queries schools, hospitals, shelters, roads, mosques, markets,
and government buildings within the Cox's Bazar bounding box.
Merges DDM shelter data when available.
"""

import requests
import geopandas as gpd
import pandas as pd
import time
import logging
from shapely.geometry import Point, LineString
from src.config import (
    COX_BAZAR_BBOX, PROCESSED_DIR, SHELTER_DATA_FILE,
    PROCESSED_SCHOOLS, PROCESSED_HOSPITALS, PROCESSED_SHELTERS,
    PROCESSED_ROADS, PROCESSED_MOSQUES, PROCESSED_GOVT_BUILDINGS,
    PROCESSED_MARKETS,
)

logger = logging.getLogger(__name__)

OVERPASS_URL = "https://overpass-api.de/api/interpreter"

QUERIES = {
    "schools": {
        "tags": '["amenity"="school"]',
        "output_file": PROCESSED_SCHOOLS,
        "geom": "point",
    },
    "hospitals": {
        "tags": '["amenity"~"hospital|clinic"]',
        "output_file": PROCESSED_HOSPITALS,
        "geom": "point",
    },
    "shelters": {
        "tags": '["amenity"="shelter"]',
        "output_file": PROCESSED_SHELTERS,
        "geom": "point",
    },
    "mosques": {
        "tags": '["amenity"="place_of_worship"]["religion"="muslim"]',
        "output_file": PROCESSED_MOSQUES,
        "geom": "point",
    },
    "govt_buildings": {
        "tags": '["office"="government"]',
        "output_file": PROCESSED_GOVT_BUILDINGS,
        "geom": "point",
    },
    "markets": {
        "tags": '["amenity"~"marketplace|market"]',
        "output_file": PROCESSED_MARKETS,
        "geom": "point",
    },
    "roads": {
        "tags": '["highway"~"primary|secondary|tertiary"]',
        "output_file": PROCESSED_ROADS,
        "geom": "line",
    },
}


def extract_all_osm_data(force=False):
    """
    Extract all infrastructure categories from OSM.

    Args:
        force: If True, re-download even if files already exist.

    Returns:
        dict of {category: GeoDataFrame}
    """
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    results = {}

    for category, cfg in QUERIES.items():
        out = cfg["output_file"]

        if out.exists() and not force:
            logger.info(f"  {category}: cached — skipping (force=True to re-extract)")
            try:
                results[category] = gpd.read_file(out)
            except Exception:
                results[category] = gpd.GeoDataFrame()
            continue

        logger.info(f"  Extracting {category} from OSM...")
        try:
            gdf = _query_overpass(cfg["tags"], cfg["geom"])
            if gdf is not None and len(gdf) > 0:
                gdf.to_file(out, driver="GeoJSON")
                results[category] = gdf
                logger.info(f"  {category}: {len(gdf)} features saved")
            else:
                logger.warning(f"  {category}: no features found")
                results[category] = gpd.GeoDataFrame()
            time.sleep(5)  # polite delay for the Overpass API
        except Exception as e:
            logger.error(f"  {category}: extraction failed — {e}")
            results[category] = gpd.GeoDataFrame()

    # Merge DDM shelter data if available
    _merge_ddm_shelters(results)

    logger.info("\n  OSM Extraction Summary:")
    for cat, gdf in results.items():
        logger.info(f"    {cat}: {len(gdf)} features")

    return results


# ── Overpass query execution ──────────────────────────────────────────────────

def _query_overpass(tags, geom_type):
    """Execute an Overpass API query and return a GeoDataFrame."""
    south, west, north, east = COX_BAZAR_BBOX
    bbox = f"{south},{west},{north},{east}"

    if geom_type == "line":
        query = f"""
        [out:json][timeout:120];
        (way{tags}({bbox}););
        out geom;
        """
    else:
        query = f"""
        [out:json][timeout:120];
        (
          node{tags}({bbox});
          way{tags}({bbox});
          relation{tags}({bbox});
        );
        out center;
        """

    resp = requests.post(OVERPASS_URL, data={"data": query}, timeout=180)
    resp.raise_for_status()
    elements = resp.json().get("elements", [])
    if not elements:
        return None

    features = []
    for el in elements:
        props = el.get("tags", {})
        props["osm_id"] = el.get("id")
        props["osm_type"] = el.get("type")

        if geom_type == "line" and el["type"] == "way":
            coords = [(n["lon"], n["lat"]) for n in el.get("geometry", [])]
            if len(coords) >= 2:
                features.append({"geometry": LineString(coords), **props})
        else:
            lat = el.get("lat") or (el.get("center") or {}).get("lat")
            lon = el.get("lon") or (el.get("center") or {}).get("lon")
            if lat and lon:
                features.append({"geometry": Point(lon, lat), **props})

    if not features:
        return None

    gdf = gpd.GeoDataFrame(features, crs="EPSG:4326")

    # Keep only useful columns
    keep = [
        "geometry", "osm_id", "osm_type", "name", "name:bn",
        "amenity", "building", "highway", "capacity", "operator",
        "addr:district", "addr:subdistrict",
    ]
    existing = [c for c in keep if c in gdf.columns]
    return gdf[existing]


# ── Merge DDM shelter data ────────────────────────────────────────────────────

def _merge_ddm_shelters(results):
    """Merge DDM shelter CSV metadata (capacity, type) into processed data.
    
    The DDM CSV has columns: Upazila, Shelter_Type, Name, Union, Capacity, Remarks
    but NO lat/lon coordinates. We store it as a reference lookup table.
    """
    if not SHELTER_DATA_FILE.exists():
        logger.info("  No DDM shelter file found — using OSM data only")
        return

    try:
        ddm = pd.read_csv(SHELTER_DATA_FILE)
        logger.info(f"  Loading DDM shelter data: {len(ddm)} records")
        logger.info(f"  DDM columns: {list(ddm.columns)}")

        # Save processed DDM data as a reference CSV
        ddm_out = PROCESSED_DIR / "ddm_shelters_reference.csv"
        ddm.to_csv(ddm_out, index=False)

        # Extract capacity statistics
        cap_col = _find_col(ddm.columns, ["capacity", "cap"])
        if cap_col:
            caps = pd.to_numeric(ddm[cap_col], errors="coerce").dropna()
            if len(caps) > 0:
                logger.info(f"  DDM shelters: {len(ddm)} total")
                logger.info(f"  Capacity range: {caps.min():.0f} – {caps.max():.0f}")
                logger.info(f"  Avg capacity: {caps.mean():.0f}")
                logger.info(f"  Total capacity: {caps.sum():.0f}")

        # List unique upazilas in DDM data
        upazila_col = _find_col(ddm.columns, ["upazila"])
        if upazila_col:
            ups = sorted(ddm[upazila_col].dropna().unique())
            logger.info(f"  DDM upazilas ({len(ups)}): {ups}")

        logger.info(f"  DDM shelter reference saved → {ddm_out}")
    except Exception as e:
        logger.error(f"  Error loading DDM shelter data: {e}")


def _find_col(columns, patterns):
    """Find a column whose lowered name contains any of the patterns."""
    for col in columns:
        cl = col.lower()
        for p in patterns:
            if p in cl:
                return col
    return None


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    extract_all_osm_data()