# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/07_distance_zonal_stats.ipynb (unless otherwise specified).

__all__ = ["create_distance_zonal_stats"]


# Cell
from typing import Any, Dict, List

import geopandas as gpd

import geowrangler.area_zonal_stats as azs
import geowrangler.vector_zonal_stats as vzs
from .vector_zonal_stats import GEO_INDEX_NAME

# Internal Cell
INTERNAL_DISTANCE_COL = "__Geowrangleer_nearest_distance"


def build_agg_distance_dicts(expanded_aggs, distance_col):
    agg_dicts = vzs._build_agg_args(expanded_aggs)
    if distance_col is not None:
        agg_dicts[distance_col] = (INTERNAL_DISTANCE_COL, "mean")
    return agg_dicts


# Cell
def create_distance_zonal_stats(
    aoi: gpd.GeoDataFrame,  # Area of interest for which zonal stats are to be computed for
    data: gpd.GeoDataFrame,  # Source gdf of region/areas containing data to compute zonal stats from
    max_distance: float,  # max distance to compute distance for (the larger the slower the join), set to None for no limit
    aggregations: List[Dict[str, Any]] = [],  # aggregations
    distance_col: str = "nearest",  # column name of the distance column, set to None if not wanted in results
):
    """Computes zonal stats based on nearest matching data geometry within `max_distance`.
    Note that setting a too high max_distance (or None) will incur a performance cost.
    """
    # aoi/data crs should be planar
    azs.validate_area_aoi(aoi)
    azs.validate_area_data(data)

    fixed_aggs = [vzs._fix_agg(agg) for agg in aggregations]

    vzs._validate_aggs(fixed_aggs, data)

    # reindex aoi
    aoi_index_name = aoi.index.name
    aoi = vzs._prep_aoi(aoi)

    # sync aoi/data crs
    if not data.crs.equals(aoi.crs):
        data = data.to_crs(aoi.crs)

    # add spatial indexes
    aoi.geometry.sindex
    data.geometry.sindex

    # nearest
    nearest = gpd.sjoin_nearest(
        aoi[[GEO_INDEX_NAME, "geometry"]],
        data,
        how="inner",
        max_distance=max_distance,
        distance_col=INTERNAL_DISTANCE_COL,
    )

    groups = nearest.groupby(GEO_INDEX_NAME)

    expanded_aggs = vzs._expand_aggs(fixed_aggs)
    agg_distance_dicts = build_agg_distance_dicts(expanded_aggs, distance_col)

    aggregates = groups.agg(**agg_distance_dicts)

    results = aoi.merge(
        aggregates, how="left", on=GEO_INDEX_NAME, suffixes=(None, "_y")
    )
    vzs._fillnas(expanded_aggs, results, aoi)

    results.set_index(GEO_INDEX_NAME, inplace=True)
    results.index.name = aoi_index_name
    return results
