geo/geo_app/cluster_index.py

"""
Server-side map clustering using Uber H3 hexagonal grid.

Maps zoom levels to h3 resolutions and groups nodes by cell.
"""
import logging
import h3

logger = logging.getLogger(__name__)

# Map zoom level to h3 resolution
# Higher zoom = higher resolution = smaller cells
ZOOM_TO_RES = {
    0: 0, 1: 0, 2: 1, 3: 1, 4: 2, 5: 2,
    6: 3, 7: 3, 8: 4, 9: 4, 10: 5, 11: 5,
    12: 6, 13: 7, 14: 8, 15: 9, 16: 10
}


def _fetch_nodes(db, west, south, east, north, transport_type=None, node_type=None):
    """Fetch nodes from database for a bounding box.

    Args:
        db: Database connection
        west, south, east, north: Bounding box coordinates
        transport_type: Filter by transport type (auto, rail, sea, air)
        node_type: Type of nodes to fetch ('logistics', 'offer', 'supplier')
    """
    bind_vars = {
        'west': west,
        'south': south,
        'east': east,
        'north': north,
    }

    # Select AQL query based on node_type
    if node_type == 'offer':
        aql = """
        FOR node IN nodes
            FILTER node.node_type == 'offer'
            FILTER node.latitude != null AND node.longitude != null
            FILTER node.latitude >= @south AND node.latitude <= @north
            FILTER node.longitude >= @west AND node.longitude <= @east
            RETURN node
        """
    elif node_type == 'supplier':
        # Get suppliers that have offers (aggregate through offers)
        aql = """
        FOR offer IN nodes
            FILTER offer.node_type == 'offer'
            FILTER offer.supplier_uuid != null
            LET supplier = DOCUMENT(CONCAT('nodes/', offer.supplier_uuid))
            FILTER supplier != null
            FILTER supplier.latitude != null AND supplier.longitude != null
            FILTER supplier.latitude >= @south AND supplier.latitude <= @north
            FILTER supplier.longitude >= @west AND supplier.longitude <= @east
            COLLECT sup_uuid = offer.supplier_uuid INTO offers
            LET sup = DOCUMENT(CONCAT('nodes/', sup_uuid))
            RETURN {
                _key: sup_uuid,
                name: sup.name,
                latitude: sup.latitude,
                longitude: sup.longitude,
                country: sup.country,
                country_code: sup.country_code,
                node_type: 'supplier',
                offers_count: LENGTH(offers)
            }
        """
    else:  # logistics (default)
        aql = """
        FOR node IN nodes
            FILTER node.node_type == 'logistics' OR node.node_type == null
            FILTER node.latitude != null AND node.longitude != null
            FILTER node.latitude >= @south AND node.latitude <= @north
            FILTER node.longitude >= @west AND node.longitude <= @east
            RETURN node
        """

    cursor = db.aql.execute(aql, bind_vars=bind_vars)
    nodes = list(cursor)

    # Filter by transport type if specified (only for logistics nodes)
    if node_type in (None, 'logistics'):
        if transport_type:
            nodes = [
                n for n in nodes
                if transport_type in (n.get('transport_types') or [])
            ]
        else:
            # Default: only rail/sea hubs
            nodes = [
                n for n in nodes
                if ('rail' in (n.get('transport_types') or [])) or ('sea' in (n.get('transport_types') or []))
            ]

    return nodes


def get_clustered_nodes(db, west, south, east, north, zoom, transport_type=None, node_type=None):
    """
    Get clustered nodes for given bounding box and zoom level.

    Uses H3 hexagonal grid to group nearby nodes.

    Args:
        db: Database connection
        west, south, east, north: Bounding box coordinates
        zoom: Map zoom level
        transport_type: Filter by transport type (for logistics nodes)
        node_type: Type of nodes ('logistics', 'offer', 'supplier')
    """
    resolution = ZOOM_TO_RES.get(int(zoom), 5)
    nodes = _fetch_nodes(db, west, south, east, north, transport_type, node_type)

    if not nodes:
        return []

    # Group nodes by h3 cell
    cells = {}
    for node in nodes:
        lat = node.get('latitude')
        lng = node.get('longitude')

        cell = h3.latlng_to_cell(lat, lng, resolution)
        if cell not in cells:
            cells[cell] = []
        cells[cell].append(node)

    # Build results
    results = []
    for cell, nodes_in_cell in cells.items():
        count = len(nodes_in_cell)

        if count == 1:
            # Single point — return actual node data
            node = nodes_in_cell[0]
            results.append({
                'id': node.get('_key'),
                'latitude': node.get('latitude'),
                'longitude': node.get('longitude'),
                'count': 1,
                'expansion_zoom': None,
                'name': node.get('name'),
            })
        else:
            # Cluster — return cell centroid
            lat, lng = h3.cell_to_latlng(cell)
            results.append({
                'id': f"cluster-{cell}",
                'latitude': lat,
                'longitude': lng,
                'count': count,
                'expansion_zoom': min(zoom + 2, 16),
                'name': None,
            })

    logger.info("Returning %d clusters/points for zoom=%d res=%d", len(results), zoom, resolution)
    return results