src/data_processing/load_boundaries.py

"""
Load and process administrative boundaries for Cox's Bazar district.

The geoBoundaries ADM4 file has a flat structure (shapeName only, no parent
hierarchy columns). We use spatial filtering with the Cox's Bazar bounding
box to extract the relevant unions.
"""

import geopandas as gpd
import logging
from shapely.geometry import box
from src.config import (
    ADMIN_BOUNDARIES_FILE, PROCESSED_UNIONS, PROCESSED_DIR, COX_BAZAR_BBOX
)

logger = logging.getLogger(__name__)

# Known Cox's Bazar upazila-level bounding boxes for spatial classification
# These help assign upazila names to unions that lack hierarchy info
UPAZILA_BOXES = {
    "Cox's Bazar Sadar": (21.35, 91.90, 21.60, 92.15),
    "Ramu":              (21.25, 92.05, 21.55, 92.30),
    "Chakaria":          (21.55, 91.90, 21.95, 92.20),
    "Kutubdia":          (21.70, 91.80, 21.95, 91.92),
    "Maheshkhali":       (21.50, 91.80, 21.80, 91.95),
    "Pekua":             (21.75, 91.92, 21.95, 92.10),
    "Teknaf":            (20.55, 92.15, 21.20, 92.40),
    "Ukhia":             (21.10, 92.05, 21.35, 92.25),
}


def load_and_filter_boundaries():
    """
    Load admin boundaries GeoJSON, spatially filter to Cox's Bazar district,
    assign upazila names based on centroid location, and save processed output.

    Returns:
        GeoDataFrame of Cox's Bazar unions.
    """
    logger.info("Loading administrative boundaries...")

    if not ADMIN_BOUNDARIES_FILE.exists():
        raise FileNotFoundError(
            f"Admin boundaries file not found: {ADMIN_BOUNDARIES_FILE}\n"
            "Download from: https://data.humdata.org/dataset/"
            "geoboundaries-admin-boundaries-for-bangladesh"
        )

    gdf = gpd.read_file(ADMIN_BOUNDARIES_FILE)
    logger.info(f"Loaded {len(gdf)} total administrative units")
    logger.info(f"Columns: {list(gdf.columns)}")

    # Ensure WGS84
    if gdf.crs is None or gdf.crs.to_epsg() != 4326:
        gdf = gdf.to_crs(epsg=4326)

    # ── Spatial filter to Cox's Bazar bbox ────────────────────────────────
    south, west, north, east = COX_BAZAR_BBOX
    cox_box = box(west, south, east, north)
    unions = gdf[gdf.geometry.intersects(cox_box)].copy()

    if len(unions) == 0:
        raise ValueError(
            f"No unions found within Cox's Bazar bbox {COX_BAZAR_BBOX}. "
            f"Dataset bounds: {gdf.total_bounds}"
        )

    logger.info(f"Found {len(unions)} unions in Cox's Bazar bounding box")

    # ── Build standardized columns ────────────────────────────────────────
    # Rename shapeName → union_name
    if "shapeName" in unions.columns:
        unions = unions.rename(columns={"shapeName": "union_name"})
    else:
        # Fallback — use whatever name column exists
        for col in unions.columns:
            if "name" in col.lower() and col != "geometry":
                unions = unions.rename(columns={col: "union_name"})
                break

    # Assign upazila names from centroid location
    centroids = unions.geometry.centroid
    upazila_names = []
    for cent in centroids:
        assigned = "Cox's Bazar (unclassified)"
        for upazila, (s, w, n, e) in UPAZILA_BOXES.items():
            if w <= cent.x <= e and s <= cent.y <= n:
                assigned = upazila
                break
        upazila_names.append(assigned)
    unions["upazila_name"] = upazila_names
    unions["district_name"] = "Cox's Bazar"

    # Keep useful columns
    keep = ["union_name", "upazila_name", "district_name", "geometry"]
    for col in ["shapeID", "shapeISO"]:
        if col in unions.columns:
            keep.append(col)
    unions = unions[[c for c in keep if c in unions.columns]].copy()

    # ── Save ──────────────────────────────────────────────────────────────
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    unions.to_file(PROCESSED_UNIONS, driver="GeoJSON")
    logger.info(f"Saved {len(unions)} unions → {PROCESSED_UNIONS}")

    # Summary
    ups = sorted(unions["upazila_name"].unique())
    logger.info(f"Upazilas ({len(ups)}): {ups}")
    for up in ups:
        n = len(unions[unions["upazila_name"] == up])
        logger.info(f"  {up}: {n} unions")

    return unions


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