Processing Module¶

def add_area_in_meters(
    gdf: gpd.GeoDataFrame, area_column_name: str = "area_in_meters"
) -> gpd.GeoDataFrame:
    """
    Calculate the area of geometries in square meters and add it as a new column.

    Automatically handles UTM transformation for accurate area calculation.

    Args:
        gdf: GeoDataFrame containing (Multi)Polygon geometries.
        area_column_name: Name of the new column.

    Returns:
        The input GeoDataFrame with an additional area column.

    Raises:
        ValueError: If the input GeoDataFrame contains non-polygon geometries.
    """
    # Validate input geometries
    if not all(gdf.geometry.geom_type.isin(["Polygon", "MultiPolygon"])):
        raise ValueError(
            "Input GeoDataFrame must contain only Polygon or MultiPolygon geometries."
        )

    # Create a copy of the GeoDataFrame to avoid modifying the original
    gdf_with_area = gdf.copy()

    # Calculate the UTM CRS for accurate area calculation
    try:
        utm_crs = gdf_with_area.estimate_utm_crs()
    except Exception as e:
        LOGGER.warning(
            f"Warning: UTM CRS estimation failed, using Web Mercator. Error: {e}"
        )
        utm_crs = "EPSG:3857"  # Fallback to Web Mercator

    # Transform to UTM CRS and calculate the area in square meters
    gdf_with_area[area_column_name] = gdf_with_area.to_crs(utm_crs).geometry.area

    return gdf_with_area

def add_spatial_jitter(
    df: pd.DataFrame,
    columns: List[str] = ["latitude", "longitude"],
    amount: float = 0.0001,
    seed=None,
    copy=True,
) -> pd.DataFrame:
    """
    Add random jitter to duplicated coordinates to separate overlapping points.

    Args:
        df: DataFrame containing coordinates.
        columns: Column names to jitter.
        amount: Amount of jitter to add.
        seed: Random seed for reproducibility.
        copy: Whether to create a copy of the input DataFrame.

    Returns:
        DataFrame with jittered coordinates.

    Raises:
        ValueError: If columns don't exist or amount is invalid.
        TypeError: If input types are incorrect.
    """

    # Input validation
    if not isinstance(df, pd.DataFrame):
        raise TypeError("Input must be a pandas DataFrame")

    if not all(col in df.columns for col in columns):
        raise ValueError(f"Not all columns {columns} found in DataFrame")

    # Handle jitter amounts
    if isinstance(amount, (int, float)):
        if amount <= 0:
            raise ValueError("Jitter amount must be positive")
        jitter_amounts = {col: amount for col in columns}
    elif isinstance(amount, dict):
        if not all(col in amount for col in columns):
            raise ValueError("Must specify jitter amount for each column")
        if not all(amt > 0 for amt in amount.values()):
            raise ValueError("All jitter amounts must be positive")
        jitter_amounts = amount
    else:
        raise TypeError("amount must be a number or dictionary")

    # Create copy if requested
    df_work = df.copy() if copy else df

    # Set random seed if provided
    if seed is not None:
        np.random.seed(seed)

    try:
        # Find duplicated coordinates
        duplicate_mask = df_work.duplicated(subset=columns, keep=False)
        n_duplicates = duplicate_mask.sum()

        if n_duplicates > 0:
            # Add jitter to each column separately
            for col in columns:
                jitter = np.random.uniform(
                    low=-jitter_amounts[col],
                    high=jitter_amounts[col],
                    size=n_duplicates,
                )
                df_work.loc[duplicate_mask, col] += jitter

            # Validate results (ensure no remaining duplicates)
            if df_work.duplicated(subset=columns, keep=False).any():
                # If duplicates remain, recursively add more jitter
                df_work = add_spatial_jitter(
                    df_work,
                    columns=columns,
                    amount={col: amt * 2 for col, amt in jitter_amounts.items()},
                    seed=seed,
                    copy=False,
                )

        return df_work

    except Exception as e:
        raise RuntimeError(f"Error during jittering operation: {str(e)}")

def aggregate_points_to_zones(
    points: Union[pd.DataFrame, gpd.GeoDataFrame],
    zones: gpd.GeoDataFrame,
    value_columns: Optional[Union[str, List[str]]] = None,
    aggregation: Union[str, Dict[str, str]] = "count",
    point_zone_predicate: str = "within",
    zone_id_column: str = "zone_id",
    output_suffix: str = "",
    drop_geometry: bool = False,
) -> gpd.GeoDataFrame:
    """
    Aggregate point data to zones with flexible aggregation methods.

    For zones with no overlapping points:
    - ``"count"`` aggregation fills missing values with ``0``.
    - All other aggregations fill missing values with ``np.nan``.

    Args:
        points: Point data to aggregate.
        zones: Zones to aggregate points to.
        value_columns: Column(s) containing values to aggregate.
        aggregation: Aggregation method(s) to use.
        point_zone_predicate: Spatial predicate (e.g., 'within', 'intersects').
        zone_id_column: Column in zones containing zone identifiers.
        output_suffix: Suffix to add to output column names.
        drop_geometry: Whether to drop the geometry column from output.

    Returns:
        Zones with aggregated point values.

    Raises:
        TypeError: If zones is not a GeoDataFrame or aggregation is invalid.
        ValueError: If columns are missing or metadata is inconsistent.
    """
    # --- Input validation ---
    if not isinstance(zones, gpd.GeoDataFrame):
        raise TypeError("zones must be a GeoDataFrame")

    if zone_id_column not in zones.columns:
        raise ValueError(f"Zone ID column '{zone_id_column}' not found in zones")

    # --- Normalise points ---
    points_gdf = (
        convert_to_geodataframe(points)
        if not isinstance(points, gpd.GeoDataFrame)
        else points.copy()
    )

    # --- CRS alignment ---
    if points_gdf.crs != zones.crs:
        points_gdf = points_gdf.to_crs(zones.crs)

    # --- Normalise value_columns ---
    if isinstance(value_columns, str):
        value_columns = [value_columns]

    if value_columns is not None:
        missing_cols = [col for col in value_columns if col not in points_gdf.columns]
        if missing_cols:
            raise ValueError(f"Value columns not found in points data: {missing_cols}")

    # --- Build agg_funcs and per-output-column method lookup ---
    agg_funcs: Dict[str, str] = {}
    # Maps output column name → aggregation method for fill-value decisions
    agg_method_for_output_col: Dict[str, str] = {}

    if isinstance(aggregation, str):
        if aggregation == "count":
            agg_funcs["__count"] = "count"
        elif value_columns is not None:
            agg_funcs = {col: aggregation for col in value_columns}
            agg_method_for_output_col = {
                f"{col}{output_suffix}": aggregation for col in value_columns
            }
        else:
            raise ValueError(
                "value_columns must be specified for aggregation methods other than 'count'"
            )
    elif isinstance(aggregation, dict):
        if value_columns is None:
            raise ValueError(
                "value_columns must be specified when using a dict of aggregation methods"
            )
        missing_aggs = [col for col in value_columns if col not in aggregation]
        extra_aggs = [col for col in aggregation if col not in value_columns]
        if missing_aggs:
            raise ValueError(f"Missing aggregation methods for columns: {missing_aggs}")
        if extra_aggs:
            raise ValueError(
                f"Aggregation methods specified for columns not in value_columns: {extra_aggs}"
            )
        agg_funcs = dict(aggregation)
        agg_method_for_output_col = {
            f"{col}{output_suffix}": method for col, method in aggregation.items()
        }
    else:
        raise TypeError("aggregation must be a str or dict")

    # --- Spatial join ---
    result = zones.copy()
    joined = gpd.sjoin(points_gdf, zones, how="inner", predicate=point_zone_predicate)

    # --- Aggregation ---
    if "__count" in agg_funcs:
        counts = (
            joined.groupby(zone_id_column)
            .size()
            .reset_index(name=f"point_count{output_suffix}")
        )
        result = result.merge(counts, on=zone_id_column, how="left")
        result[f"point_count{output_suffix}"] = (
            result[f"point_count{output_suffix}"].fillna(0).astype(int)
        )
    else:
        # Drop geometry before non-count aggregations to avoid errors
        if "geometry" in joined.columns:
            joined = joined.drop(columns=["geometry"])

        aggregated = joined.groupby(zone_id_column).agg(agg_funcs).reset_index()

        # Flatten MultiIndex columns produced by some pandas agg paths
        if isinstance(aggregated.columns, pd.MultiIndex):
            aggregated.columns = [
                (
                    f"{col[0]}_{col[1]}{output_suffix}"
                    if col[0] != zone_id_column
                    else zone_id_column
                )
                for col in aggregated.columns
            ]
        else:
            # Single-level: rename value columns to include suffix
            aggregated = aggregated.rename(
                columns={
                    col: f"{col}{output_suffix}"
                    for col in aggregated.columns
                    if col != zone_id_column
                }
            )

        result = result.merge(aggregated, on=zone_id_column, how="left")

        # -------------------------------------------------------
        # Fill with 0 only for 'count', NaN for everything
        # else so zones with no overlapping points are distinguishable
        # from zones whose true aggregated value is zero.
        # -------------------------------------------------------
        for col in result.columns:
            if col in (zone_id_column, "geometry"):
                continue
            if not pd.api.types.is_numeric_dtype(result[col]):
                continue
            method = agg_method_for_output_col.get(col, "")
            fill_value = 0 if method == "count" else np.nan
            result[col] = result[col].fillna(fill_value)

    if drop_geometry:
        result = result.drop(columns=["geometry"])

    return result

def aggregate_polygons_to_zones(
    polygons: Union[pd.DataFrame, gpd.GeoDataFrame],
    zones: gpd.GeoDataFrame,
    value_columns: Union[str, List[str]],
    aggregation: Union[str, Dict[str, str]] = "sum",
    predicate: Literal["intersects", "within", "fractional"] = "intersects",
    zone_id_column: str = "zone_id",
    output_suffix: str = "",
    drop_geometry: bool = False,
) -> gpd.GeoDataFrame:
    """
    Aggregates polygon data to zones based on a spatial relationship.

    Args:
        polygons: Polygon data to aggregate.
        zones: Target zones.
        value_columns: Column(s) in polygons with numeric values.
        aggregation: Aggregation method(s) to use (str or dict).
        predicate: Spatial relationship ('intersects', 'within', 'fractional').
        zone_id_column: Unique identifier column in zones.
        output_suffix: Suffix for output columns.
        drop_geometry: Whether to drop the geometry column.

    Returns:
        The zones GeoDataFrame with aggregated values.

    Raises:
        TypeError: If polygons/zones are not GeoDataFrames.
        ValueError: If columns or predicates are invalid.
    """
    # --- Input validation ---
    if not isinstance(zones, gpd.GeoDataFrame):
        raise TypeError("zones must be a GeoDataFrame")

    if zones.empty:
        raise ValueError("zones GeoDataFrame is empty")

    if zone_id_column not in zones.columns:
        raise ValueError(f"Zone ID column '{zone_id_column}' not found in zones")

    if predicate not in ["intersects", "within", "fractional"]:
        raise ValueError(
            f"Unsupported predicate: '{predicate}'. "
            "Must be one of: 'intersects', 'within', 'fractional'."
        )

    # --- Normalise polygons ---
    if not isinstance(polygons, gpd.GeoDataFrame):
        try:
            polygons_gdf = convert_to_geodataframe(polygons)
        except Exception as e:
            raise TypeError(
                f"polygons must be a GeoDataFrame or convertible to one: {e}"
            )
    else:
        polygons_gdf = polygons.copy()

    if polygons_gdf.empty:
        LOGGER.warning("Empty polygons GeoDataFrame provided")
        return zones

    # --- Geometry type validation ---
    non_polygon_geoms = [
        geom_type
        for geom_type in polygons_gdf.geometry.geom_type.unique()
        if geom_type not in ["Polygon", "MultiPolygon"]
    ]
    if non_polygon_geoms:
        raise ValueError(
            f"Input contains non-polygon geometries: {non_polygon_geoms}. "
            "Use aggregate_points_to_zones for point data."
        )

    # --- Normalise value_columns ---
    if isinstance(value_columns, str):
        value_columns = [value_columns]

    missing_cols = [col for col in value_columns if col not in polygons_gdf.columns]
    if missing_cols:
        raise ValueError(f"Value columns not found in polygons data: {missing_cols}")

    if zone_id_column in polygons_gdf.columns:
        raise ValueError(
            f"Column name conflict: polygons DataFrame contains column '{zone_id_column}' "
            "which conflicts with the zone identifier column. "
            "Please rename this column in the polygons data."
        )

    # --- CRS alignment ---
    if polygons_gdf.crs != zones.crs:
        polygons_gdf = polygons_gdf.to_crs(zones.crs)

    # --- Build aggregation functions ---
    agg_funcs = _process_aggregation_methods(aggregation, value_columns)

    # Build lookup: original col name → method (before suffix is applied)
    # Used below to decide fill value per column.
    if isinstance(aggregation, str):
        agg_method_for_col: Dict[str, str] = {col: aggregation for col in value_columns}
    else:
        agg_method_for_col = dict(aggregation)

    # --- Spatial aggregation ---
    minimal_zones = zones[[zone_id_column, "geometry"]].copy()

    if predicate == "fractional":
        aggregated_data = _fractional_aggregation(
            polygons_gdf, minimal_zones, value_columns, agg_funcs, zone_id_column
        )
    else:
        aggregated_data = _simple_aggregation(
            polygons_gdf,
            minimal_zones,
            value_columns,
            agg_funcs,
            zone_id_column,
            predicate,
        )

    # --- Merge back to full zones ---
    result = zones.merge(
        aggregated_data[[col for col in aggregated_data.columns if col != "geometry"]],
        on=zone_id_column,
        how="left",
    )

    # --- Fill NaN values: 0 for count, np.nan for everything else ---
    # NOTE: output_suffix has NOT been applied yet, so column names here
    # still match the keys in agg_method_for_col.
    aggregated_cols = [col for col in result.columns if col not in zones.columns]
    for col in aggregated_cols:
        if not pd.api.types.is_numeric_dtype(result[col]):
            continue
        method = agg_method_for_col.get(col, "")
        fill_value = 0 if method == "count" else np.nan
        result[col] = result[col].fillna(fill_value)

    # --- Apply output suffix ---
    if output_suffix:
        rename_dict = {col: f"{col}{output_suffix}" for col in aggregated_cols}
        result = result.rename(columns=rename_dict)

    if drop_geometry:
        result = result.drop(columns=["geometry"])

    return result

def annotate_with_admin_regions(
    gdf: gpd.GeoDataFrame,
    country_code: str,
    data_store: Optional[DataStore] = None,
    admin_id_column_suffix="",
) -> gpd.GeoDataFrame:
    """
    Annotate a GeoDataFrame with administrative region information.

    Performs a spatial join between the input points and administrative boundaries
    at levels 1 and 2, resolving conflicts when points intersect multiple regions.

    Args:
        gdf: GeoDataFrame containing points to annotate.
        country_code: ISO country code for administrative boundaries.
        data_store: Optional DataStore for loading boundary data.
        admin_id_column_suffix: Optional suffix for admin ID columns.

    Returns:
        GeoDataFrame with added administrative region columns (admin1, admin2, etc.).

    Raises:
        TypeError: If gdf is not a GeoDataFrame.
    """
    from gigaspatial.handlers.boundaries import AdminBoundaries

    if not isinstance(gdf, gpd.GeoDataFrame):
        raise TypeError("gdf must be a GeoDataFrame")

    if gdf.empty:
        LOGGER.warning("Empty GeoDataFrame provided, returning as-is")
        return gdf

    # read country admin data
    admin1_data = AdminBoundaries.create(
        country_code=country_code, admin_level=1, data_store=data_store
    ).to_geodataframe()

    admin1_data.rename(
        columns={"boundary_id": f"admin1_id{admin_id_column_suffix}", "name": "admin1"},
        inplace=True,
    )
    admin1_data = admin1_data[[f"admin1_id{admin_id_column_suffix}", "admin1", "geometry"]]

    admin2_data = AdminBoundaries.create(
        country_code=country_code, admin_level=2, data_store=data_store
    ).to_geodataframe()

    admin2_data.rename(
        columns={
            "boundary_id": f"admin2_id{admin_id_column_suffix}",
            "parent_id": f"admin1_id{admin_id_column_suffix}",
            "name": "admin2",
        },
        inplace=True,
    )
    admin2_data = admin2_data[
        [
            f"admin2_id{admin_id_column_suffix}",
            "admin2",
            f"admin1_id{admin_id_column_suffix}",
            "geometry"
        ]
    ]

    # Join dataframes based on 'admin1_id_giga'
    admin_data = admin2_data.merge(
        admin1_data,
        left_on=f"admin1_id{admin_id_column_suffix}",
        right_on=f"admin1_id{admin_id_column_suffix}",
        how="outer",
    )

    admin_data["geometry"] = admin_data.apply(
        lambda x: x.geometry_x if x.geometry_x else x.geometry_y, axis=1
    )

    admin_data = gpd.GeoDataFrame(
        admin_data.drop(columns=["geometry_x", "geometry_y"]),
        geometry="geometry",
        crs=4326,
    )

    # admin_data["admin2"].fillna("Unknown", inplace=True)
    admin_data[f"admin2_id{admin_id_column_suffix}"] = admin_data[
        f"admin2_id{admin_id_column_suffix}"
    ].replace({np.nan: None})

    if gdf.crs is None:
        LOGGER.warning("Input GeoDataFrame has no CRS, assuming EPSG:4326")
        gdf.set_crs(epsg=4326, inplace=True)
    elif gdf.crs != "EPSG:4326":
        LOGGER.info(f"Reprojecting from {gdf.crs} to EPSG:4326")
        gdf = gdf.to_crs(epsg=4326)

    # spatial join gdf to admins
    gdf_w_admins = gdf.copy().sjoin(
        admin_data,
        how="left",
        predicate="intersects",
    )

    # Check for duplicates caused by points intersecting multiple polygons
    if len(gdf_w_admins) != len(gdf):
        LOGGER.warning(
            "Some points intersect multiple administrative boundaries. Resolving conflicts..."
        )

        # Group by original index and select the closest admin area for ties
        gdf_w_admins["distance"] = gdf_w_admins.apply(
            lambda row: row.geometry.distance(
                admin_data.loc[row.index_right, "geometry"].centroid
            ),
            axis=1,
        )

        # For points with multiple matches, keep the closest polygon
        gdf_w_admins = gdf_w_admins.loc[
            gdf_w_admins.groupby(gdf.index)["distance"].idxmin()
        ].drop(columns="distance")

    # Drop unnecessary columns and reset the index
    gdf_w_admins = gdf_w_admins.drop(columns="index_right").reset_index(drop=True)

    return gdf_w_admins

def buffer_geodataframe(
    gdf: gpd.GeoDataFrame,
    buffer_distance_meters: Union[float, np.array, pd.Series],
    cap_style: Literal["round", "square", "flat"] = "round",
    copy=True,
) -> gpd.GeoDataFrame:
    """
    Buffer a GeoDataFrame with a distance in meters.

    Automatically handles UTM transformation for accurate buffering.

    Args:
        gdf: The GeoDataFrame to be buffered.
        buffer_distance_meters: The buffer distance.
        cap_style: Style of caps ('round', 'square', 'flat').
        copy: Whether to create a copy of the input.

    Returns:
        The buffered GeoDataFrame in the original CRS.
    """

    # Input validation
    if not isinstance(gdf, gpd.GeoDataFrame):
        raise TypeError("Input must be a GeoDataFrame")

    if cap_style not in ["round", "square", "flat"]:
        raise ValueError("cap_style must be round, flat or square.")

    if gdf.crs is None:
        raise ValueError("Input GeoDataFrame must have a defined CRS")

    # Create a copy if requested
    gdf_work = gdf.copy() if copy else gdf

    # Store input CRS
    input_crs = gdf_work.crs

    try:
        try:
            utm_crs = gdf_work.estimate_utm_crs()
        except Exception as e:
            LOGGER.warning(
                f"Warning: UTM CRS estimation failed, using Web Mercator. Error: {e}"
            )
            utm_crs = "EPSG:3857"  # Fallback to Web Mercator

        # Transform to UTM, create buffer, and transform back
        gdf_work = gdf_work.to_crs(utm_crs)
        gdf_work["geometry"] = gdf_work["geometry"].buffer(
            distance=buffer_distance_meters, cap_style=cap_style
        )
        gdf_work = gdf_work.to_crs(input_crs)

        return gdf_work

    except Exception as e:
        raise RuntimeError(f"Error during buffering operation: {str(e)}")