dopecarpet/slice.py

import argparse
import duckdb
import geohash
import geopandas as gpd
import pandas as pd
import json
import cbor2
from pathlib import Path
from shapely import wkb
from shapely.geometry import box
from typing import List, Dict

# Configure logging
import logging
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)

def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
    """Generates all 4-character Geohashes that intersect a bounding box."""
    geohashes = set()
    step_size = 0.25  # 4-character Geohash grid size

    lat = min_y
    while lat <= max_y:
        lon = min_x
        while lon <= max_x:
            gh = geohash.encode(lat, lon, precision=4)
            geohashes.add(gh)
            lon += step_size
        lat += step_size

    return sorted(geohashes)

def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame):
    """Appends new data to an existing Parquet file while ensuring correct merging."""
    if file_path.exists():
        try:
            existing_data = gpd.read_parquet(file_path)
            combined_data = pd.concat([existing_data, new_data], ignore_index=True)

            if "feature_id" in combined_data.columns:
                combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last")

            combined_data = combined_data.reset_index(drop=True)
            combined_data.to_parquet(file_path, index=False)
            logger.info(f"✅ Appended new data to {file_path}")
        except Exception as e:
            logger.warning(f"❌ Error while merging {file_path}: {e}")
    else:
        new_data.to_parquet(file_path, index=False)
        logger.info(f"✅ Created new file: {file_path}")

def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a GeoJSON file."""
    geojson_file = output_path / f"{geohash_code}.geojson"
    geojson_dict = json.loads(gdf.to_json())

    with open(geojson_file, "w", encoding="utf-8") as f:
        json.dump(geojson_dict, f)
    logger.info(f"✅ Saved GeoJSON: {geojson_file}")

def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a CBOR file (compact binary format)."""
    cbor_file = output_path / f"{geohash_code}.cbor"
    geojson_dict = json.loads(gdf.to_json())

    with open(cbor_file, "wb") as f:
        cbor2.dump(geojson_dict, f)
    logger.info(f"✅ Saved CBOR: {cbor_file}")

def geohash_bbox(geohash_code):
    """Returns a bounding box polygon for a geohash."""
    bbox = geohash.bbox(geohash_code)
    return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"])

def ensure_crs_consistency(gdf):
    """Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches."""
    if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326":
        gdf = gdf.to_crs("EPSG:4326")
    return gdf

def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame):
    """Adds water polygons to a geohash tile."""
    logger.info(f"Processing {geohash_file}...")

    # Load existing geohash file
    try:
        existing_gdf = gpd.read_parquet(geohash_file)
    except Exception as e:
        logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}")
        return

    existing_gdf = ensure_crs_consistency(existing_gdf)

    # Get the geohash bounding box
    bbox = geohash_bbox(geohash_code)

    # Extract water polygons that overlap with this geohash
    water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy()

    if water_in_tile.empty:
        logger.info(f"No water found for {geohash_code}. Skipping water addition.")
        return

    # Clip water polygons to the geohash tile boundary
    water_in_tile["geometry"] = water_in_tile.intersection(bbox)

    # Assign "water" metadata
    water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile)

    # Ensure columns match before merging
    for col in ["feature_id", "tags"]:
        if col not in existing_gdf.columns:
            existing_gdf[col] = None
        if col not in water_in_tile.columns:
            water_in_tile[col] = None

    # Merge updated water polygons into geohash file
    updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326")

    # Save back to parquet
    updated_gdf.to_parquet(geohash_file, index=False)
    logger.info(f"Updated geohash file {geohash_file} with water polygons.")

def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str, export_parquet: bool, export_geojson: bool, export_cbor: bool):
    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and adds water polygons."""
    input_path, output_path, water_path = Path(input_file), Path(output_dir), Path(water_file)
    if not input_path.exists():
        raise FileNotFoundError(f"File '{input_file}' not found.")
    if not water_path.exists():
        raise FileNotFoundError(f"File '{water_file}' not found.")
    output_path.mkdir(parents=True, exist_ok=True)

    logger.info(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")

    # Load water dataset
    logger.info(f"Loading water dataset from {water_file}...")
    water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path))

    # Connect to DuckDB and load spatial extension
    con = duckdb.connect()
    con.execute("INSTALL spatial; LOAD spatial;")

    # Load the GeoParquet file into DuckDB
    con.execute(f"CREATE TEMP TABLE geoparquet AS SELECT * FROM read_parquet('{input_path}') WHERE ST_GeometryType(geometry) IS NOT NULL AND ST_GeometryType(geometry) != 'POINT'")

    # Precompute bounding boxes for all features
    bbox_results = con.execute("""
        SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)),
                     ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry))
        FROM geoparquet;
    """).fetchall()

    # Map geohash codes to rowids
    geohash_mapping: Dict[str, List[int]] = {}
    for feature_id, rowid, min_x, min_y, max_x, max_y in bbox_results:
        intersecting_geohashes = get_geohashes_from_bbox(min_x, min_y, max_x, max_y)
        for geohash_code in intersecting_geohashes:
            if geohash_code not in geohash_mapping:
                geohash_mapping[geohash_code] = []
            geohash_mapping[geohash_code].append(rowid)

    # Process each geohash tile
    for geohash_code, rowids in geohash_mapping.items():
        geohash_bbox = geohash.bbox(geohash_code)

        # Fetch and clip geometries for this geohash tile
        filtered_data = con.execute(f"""
            SELECT feature_id, tags, 
                   ST_AsWKB(ST_Intersection(geometry, 
                                            ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                                            {geohash_bbox["e"]}, {geohash_bbox["n"]}))) 
                   AS clipped_geom 
            FROM geoparquet 
            WHERE rowid IN ({','.join(map(str, rowids))})
            AND ST_Intersects(geometry, 
                              ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                              {geohash_bbox["e"]}, {geohash_bbox["n"]}));
        """).fetchdf()

        if not filtered_data.empty:
            # Convert WKB geometries to Shapely geometries
            filtered_data["geometry"] = filtered_data["clipped_geom"].apply(
                lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None
            )
            filtered_data.drop(columns=["clipped_geom"], inplace=True)

            # Drop rows with invalid geometries
            filtered_data = filtered_data.dropna(subset=["geometry"])

            if not filtered_data.empty:
                gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")

                # Export to desired formats
                if export_parquet:
                    geohash_file = output_path / f"{geohash_code}.parquet"
                    append_to_parquet(geohash_file, gdf)
                    add_water_to_geohash(geohash_code, geohash_file, water_gdf)
                
                if export_geojson:
                    save_as_geojson(geohash_code, gdf, output_path)
                
                if export_cbor:
                    save_as_cbor(geohash_code, gdf, output_path)

    con.close()
    logger.info("🎉 Processing complete!")

def main():
    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, add water polygons, and export in desired formats.")
    parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.")
    parser.add_argument("-o", "--output", required=True, help="Directory to save output files.")
    parser.add_argument("-w", "--water", required=True, help="Path to the water.parquet file.")
    parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet")
    parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON")
    parser.add_argument("--cbor", action="store_true", help="Export as CBOR")
    
    args = parser.parse_args()

    if not (args.parquet or args.geojson or args.cbor):
        args.parquet = True  # Default to GeoParquet if no options are given

    try:
        slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor)
    except Exception as e:
        logger.error(f"❌ Error: {e}")

if __name__ == "__main__":
    main()