"""
H3 spatial indexing for geo-filtered satellite image retrieval.

Uses Uber's H3 hexagonal grid system for efficient spatial queries.
"""

import h3
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass


@dataclass
class GeoBox:
    """Bounding box for spatial queries."""
    lat_min: float
    lat_max: float
    lon_min: float
    lon_max: float


class SpatialIndex:
    """H3-based spatial index for satellite images."""
    
    def __init__(self, resolution: int = 7):
        """
        Initialize spatial index.
        
        Args:
            resolution: H3 resolution (0-15). Level 7 ≈ 1.2km cells.
        """
        self.resolution = resolution
        self.h3_to_indices: Dict[str, List[int]] = {}
        self.index_to_geo: Dict[int, Tuple[float, float]] = {}
    
    def add_image(self, index: int, lat: float, lon: float) -> str:
        """Add image coordinates to spatial index. Returns H3 cell."""
        h3_cell = h3.latlng_to_cell(lat, lon, self.resolution)
        
        if h3_cell not in self.h3_to_indices:
            self.h3_to_indices[h3_cell] = []
        self.h3_to_indices[h3_cell].append(index)
        self.index_to_geo[index] = (lat, lon)
        
        return h3_cell
    
    def query_radius(self, lat: float, lon: float, radius_km: float = 10.0) -> List[int]:
        """
        Find all images within radius of a point.
        
        Args:
            lat: Center latitude
            lon: Center longitude
            radius_km: Search radius in kilometers
            
        Returns:
            List of image indices within radius
        """
        # Get H3 cells within radius
        center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
        ring = h3.grid_disk(center_cell, k=max(1, int(radius_km / 5)))
        
        # Collect all image indices in matching cells
        results = []
        for cell in ring:
            if cell in self.h3_to_indices:
                results.extend(self.h3_to_indices[cell])
        
        return results
    
    def query_bbox(self, bbox: GeoBox) -> List[int]:
        """
        Find all images within a bounding box.
        
        Args:
            bbox: Bounding box with lat/lon bounds
            
        Returns:
            List of image indices within bbox
        """
        results = []
        for index, (lat, lon) in self.index_to_geo.items():
            if (bbox.lat_min <= lat <= bbox.lat_max and 
                bbox.lon_min <= lon <= bbox.lon_max):
                results.append(index)
        return results
    
    def get_neighbors(self, index: int, k: int = 1) -> List[int]:
        """Get neighboring image indices."""
        if index not in self.index_to_geo:
            return []
        
        lat, lon = self.index_to_geo[index]
        center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
        ring = h3.grid_disk(center_cell, k=k)
        
        results = []
        for cell in ring:
            if cell in self.h3_to_indices:
                results.extend(self.h3_to_indices[cell])
        
        return [i for i in results if i != index]
    
    def get_cell_count(self) -> int:
        """Number of occupied H3 cells."""
        return len(self.h3_to_indices)
    
    def get_image_count(self) -> int:
        """Total number of indexed images."""
        return len(self.index_to_geo)


def lat_lon_to_h3(lat: float, lon: float, resolution: int = 7) -> str:
    """Convert lat/lon to H3 cell index."""
    return h3.latlng_to_cell(lat, lon, resolution)


def h3_to_lat_lon(h3_cell: str) -> Tuple[float, float]:
    """Convert H3 cell to center lat/lon."""
    return h3.cell_to_latlng(h3_cell)


def get_h3_neighbors(h3_cell: str, k: int = 1) -> List[str]:
    """Get neighboring H3 cells."""
    return list(h3.grid_disk(h3_cell, k=k))