"""Aggregation helpers."""

import gc
import multiprocessing as mp
import os
from collections import defaultdict
from collections.abc import Callable, Generator
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field
from functools import cache
from typing import Literal

import antimeridian
import cudf
import cuml.cluster
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import odc.geo.geobox
import odc.geo.types
import pandas as pd
import psutil
import shapely
import shapely.ops
import sklearn
import xarray as xr
import xdggs
import xvec
from rich import print
from rich.progress import track
from shapely.geometry import LineString, Polygon
from stopuhr import stopwatch
from xdggs.healpix import HealpixInfo

from entropice.spatial import grids
from entropice.utils.types import Grid


@dataclass(frozen=True)
class _Aggregations:
    # ! The ordering is super important for this class!
    _common: bool = False  # If true, use _agg_cell_data_single_common and ignore orther flags
    mean: bool = True
    sum: bool = False
    std: bool = False
    min: bool = False
    max: bool = False
    median: bool = False  # Will be mapped as quantile 0.5
    quantiles: tuple[float] = field(default_factory=tuple)

    def __post_init__(self):
        assert isinstance(self.quantiles, tuple), "Quantiles must be a tuple, ortherwise the class is not hashable"
        if self.median and 0.5 in self.quantiles:
            raise ValueError("Median aggregation cannot be used together with quantile 0.5")

    @classmethod
    def common(cls):
        return cls(_common=True)

    @cache
    def __len__(self) -> int:
        if self._common:
            return 11
        length = 0
        if self.mean:
            length += 1
        if self.sum:
            length += 1
        if self.std:
            length += 1
        if self.min:
            length += 1
        if self.max:
            length += 1
        if self.median and 0.5 not in self.quantiles:
            length += 1
        length += len(self.quantiles)
        return length

    def aggnames(self) -> list[str]:
        if self._common:
            return [
                "mean",
                "std",
                "min",
                "max",
                "median",
                "p1",
                "p5",
                "p25",
                "p75",
                "p95",
                "p99",
            ]
        names = []
        if self.mean:
            names.append("mean")
        if self.sum:
            names.append("sum")
        if self.std:
            names.append("std")
        if self.min:
            names.append("min")
        if self.max:
            names.append("max")
        if self.median:
            names.append("median")
        for q in sorted(self.quantiles):  # ? Ordering is important here!
            q_int = int(q * 100)
            names.append(f"p{q_int}")
        return names

    def _agg_cell_data_single_common(self, flattened_var: xr.DataArray) -> np.ndarray:
        others_shape = tuple([flattened_var.sizes[dim] for dim in flattened_var.dims if dim != "z"])
        cell_data = np.full((len(self), *others_shape), np.nan, dtype=np.float32)
        cell_data[0, ...] = flattened_var.mean(dim="z", skipna=True).to_numpy()
        cell_data[1, ...] = flattened_var.std(dim="z", skipna=True).to_numpy()
        cell_data[2, ...] = flattened_var.min(dim="z", skipna=True).to_numpy()
        cell_data[3, ...] = flattened_var.max(dim="z", skipna=True).to_numpy()
        quantiles_to_compute = [
            0.5,
            0.01,
            0.05,
            0.25,
            0.75,
            0.95,
            0.99,
        ]  # ? Ordering is important here!
        cell_data[4:, ...] = flattened_var.quantile(
            q=quantiles_to_compute,
            dim="z",
            skipna=True,
        ).to_numpy()
        return cell_data

    def _agg_cell_data_single(self, flattened_var: xr.DataArray) -> np.ndarray:
        if self._common:
            return self._agg_cell_data_single_common(flattened_var)
        others_shape = tuple([flattened_var.sizes[dim] for dim in flattened_var.dims if dim != "z"])
        cell_data = np.full((len(self), *others_shape), np.nan, dtype=np.float32)
        j = 0
        if self.mean:
            cell_data[j, ...] = flattened_var.mean(dim="z", skipna=True).to_numpy()
            j += 1
        if self.sum:
            cell_data[j, ...] = flattened_var.sum(dim="z", skipna=True).to_numpy()
            j += 1
        if self.std:
            cell_data[j, ...] = flattened_var.std(dim="z", skipna=True).to_numpy()
            j += 1
        if self.min:
            cell_data[j, ...] = flattened_var.min(dim="z", skipna=True).to_numpy()
            j += 1
        if self.max:
            cell_data[j, ...] = flattened_var.max(dim="z", skipna=True).to_numpy()
            j += 1
        if len(self.quantiles) > 0 or self.median:
            quantiles_to_compute = sorted(self.quantiles)  # ? Ordering is important here!
            if self.median:
                quantiles_to_compute.insert(0, 0.5)

            # Potential way to use cupy:
            # cp.apply_along_axis(
            #     lambda x: cp.quantile(x[~cp.isnan(x)], q=quantiles), axis=0, #     arr=data_gpu_raw
            # )

            cell_data[j:, ...] = flattened_var.quantile(
                q=quantiles_to_compute,
                dim="z",
                skipna=True,
            ).to_numpy()
        return cell_data

    def agg_cell_data(self, flattened: xr.Dataset | xr.DataArray) -> np.ndarray:
        if isinstance(flattened, xr.DataArray):
            return self._agg_cell_data_single(flattened)

        others_shape = tuple([flattened.sizes[dim] for dim in flattened.dims if dim != "z"])
        cell_data = np.full(
            (len(flattened.data_vars), len(self), *others_shape),
            np.nan,
            dtype=np.float32,
        )
        for i, var in enumerate(flattened.data_vars):
            cell_data[i, ...] = self._agg_cell_data_single(flattened[var])
        # Transform to numpy arrays
        # if flattened.cupy.is_cupy:
        #     cell_data = cell_data.get()
        return cell_data


def _crosses_antimeridian(geom: Polygon) -> bool:
    coords = shapely.get_coordinates(geom)
    crosses_any_meridian = (coords[:, 0] > 0).any() and (coords[:, 0] < 0).any()
    return crosses_any_meridian and abs(coords[:, 0]).max() > 90


def _split_antimeridian_cell(geom: Polygon) -> list[Polygon]:
    # Assumes that it is a antimeridian hex
    coords = shapely.get_coordinates(geom)
    for i in range(coords.shape[0]):
        if coords[i, 0] < 0:
            coords[i, 0] += 360
    geom = Polygon(coords)
    antimeridian = LineString([[180, -90], [180, 90]])
    polys = shapely.ops.split(geom, antimeridian)
    return list(polys.geoms)


def _check_geom(geobox: odc.geo.geobox.GeoBox, geom: odc.geo.Geometry) -> bool:
    enclosing = geobox.enclosing(geom)
    x, y = enclosing.shape
    if x <= 1 or y <= 1:
        return False
    roi: odc.geo.types.NormalizedROI = geobox.overlap_roi(enclosing)
    roix, roiy = roi
    return (roix.stop - roix.start) > 1 and (roiy.stop - roiy.start) > 1


@stopwatch("Correcting geometries", log=False)
def _get_corrected_geoms(
    inp: tuple[Polygon, odc.geo.geobox.GeoBox, str],
) -> list[odc.geo.Geometry]:
    geom, gbox, crs = inp
    # cell.geometry is a shapely Polygon
    if crs != "EPSG:4326" or not _crosses_antimeridian(geom):
        geoms = [geom]
    # Split geometry in case it crossed antimeridian
    else:
        geoms = _split_antimeridian_cell(geom)
    geoms = [odc.geo.Geometry(g, crs=crs) for g in geoms]
    geoms = list(filter(lambda g: _check_geom(gbox, g), geoms))
    return geoms


@stopwatch("Extracting cell data", log=False)
def _extract_cell_data(cropped: xr.Dataset | xr.DataArray, aggregations: _Aggregations):
    spatdims = ["latitude", "longitude"] if "latitude" in cropped.dims and "longitude" in cropped.dims else ["y", "x"]
    flattened = cropped.stack(z=spatdims)  # noqa: PD013
    # if flattened.z.size > 1000000:
    #     flattened = flattened.cupy.as_cupy()
    cell_data = aggregations.agg_cell_data(flattened)
    return cell_data


@stopwatch("Extracting split cell data", log=False)
def _extract_split_cell_data(cropped_list: list[xr.Dataset] | list[xr.DataArray], aggregations: _Aggregations):
    spatdims = (
        ["latitude", "longitude"]
        if "latitude" in cropped_list[0].dims and "longitude" in cropped_list[0].dims
        else ["y", "x"]
    )
    flattened = xr.concat([c.stack(z=spatdims) for c in cropped_list], dim="z")  # noqa: PD013
    # if flattened.z.size > 1000000:
    #     flattened = flattened.cupy.as_cupy()
    cell_data = aggregations.agg_cell_data(flattened)
    return cell_data


@stopwatch("Cropping (and reading?) cell data", log=False)
def _read_cell_data(unaligned: xr.Dataset | xr.DataArray, geom: odc.geo.Geometry) -> xr.Dataset:
    cropped: xr.Dataset = unaligned.odc.crop(geom, apply_mask=True).drop_vars("spatial_ref")
    return cropped.compute()


@stopwatch("Cropping (and reading?) split cell data", log=False)
def _read_split_cell_data(unaligned: xr.Dataset | xr.DataArray, geoms: list[odc.geo.Geometry]) -> list[xr.Dataset]:
    cropped: list[xr.Dataset] = [unaligned.odc.crop(geom, apply_mask=True).drop_vars("spatial_ref") for geom in geoms]
    return [c.compute() for c in cropped]


@stopwatch("Processing cell", log=False)
def _process_geom(poly: Polygon, unaligned: xr.Dataset | xr.DataArray, aggregations: _Aggregations):
    geoms = _get_corrected_geoms((poly, unaligned.odc.geobox, unaligned.odc.crs))
    if len(geoms) == 0:
        raise ValueError("No geometries to process")
    elif len(geoms) == 1:
        cropped = _read_cell_data(unaligned, geoms[0])
        cell_data = _extract_cell_data(cropped, aggregations)
    else:
        cropped = _read_split_cell_data(unaligned, geoms)
        cell_data = _extract_split_cell_data(cropped, aggregations)
    return cell_data


def partition_grid(grid_gdf: gpd.GeoDataFrame, n_partitions: int, plot: bool = False) -> Generator[gpd.GeoDataFrame]:
    """Partition the input GeoDataFrame into n_partitions parts.

    Args:
        grid_gdf (gpd.GeoDataFrame): The input GeoDataFrame to partition.
        n_partitions (int): The number of partitions.
        plot (bool, optional): Whether to plot the partitions. Defaults to True.

    Yields:
        Generator[gpd.GeoDataFrame]: Partitions of the input GeoDataFrame.

    """
    if grid_gdf.crs.to_epsg() == 4326:
        crosses_antimeridian = grid_gdf.geometry.apply(_crosses_antimeridian)
    else:
        crosses_antimeridian = pd.Series([False] * len(grid_gdf), index=grid_gdf.index)

    if crosses_antimeridian.any():
        grid_gdf_am = grid_gdf[crosses_antimeridian]
        grid_gdf = grid_gdf[~crosses_antimeridian]

    # Simple partitioning by splitting the GeoDataFrame into n_partitions parts
    centroids = pd.DataFrame({"x": grid_gdf.geometry.centroid.x, "y": grid_gdf.geometry.centroid.y})

    # use cuml and cudf if len of centroids is larger than 100000
    if len(centroids) > 100000:
        print(f"Using cuML KMeans for partitioning {len(centroids)} centroids")
        centroids_cudf = cudf.DataFrame.from_pandas(centroids)
        kmeans = cuml.cluster.KMeans(n_clusters=n_partitions, random_state=42)
        labels = kmeans.fit_predict(centroids_cudf).to_pandas().to_numpy()
    else:
        labels = sklearn.cluster.KMeans(n_clusters=n_partitions, random_state=42).fit_predict(centroids)

    if plot:
        grid_gdf = grid_gdf.copy()
        grid_gdf["partition"] = labels
        ax = grid_gdf.plot(column="partition", categorical=True, legend=True, figsize=(10, 10))
        if crosses_antimeridian.any():
            grid_gdf_am.plot(ax=ax, color="red", edgecolor="black", alpha=0.5)
        ax.set_title("Grid partitions")
        plt.show()
    for i in range(n_partitions):
        partition = grid_gdf[labels == i]
        yield partition

    if crosses_antimeridian.any():
        yield grid_gdf_am


class _MemoryProfiler:
    def __init__(self):
        self.process = psutil.Process()
        self._logs = defaultdict(list)

    def log_memory(self, message: str, log: bool = True):
        mem = self.process.memory_info().rss / 1024**3
        self._logs[message].append(mem)
        if log:
            print(f"[MemoryProfiler] {message}: {mem:.2f} GB")

    def summary(self):
        summary_lines = ["[MemoryProfiler] Memory usage summary:"]
        for message, mems in self._logs.items():
            mems_str = [f"{m:.2f}" for m in mems]
            summary_lines.append(f"  {message}: min={min(mems):.2f} GB, max={max(mems):.2f} GB")
            summary_lines.append(f"    Measurements: {', '.join(mems_str)}")
        return "\n".join(summary_lines)


memprof = None
shared_raster = None


def _init_worker(r: xr.Dataset | None):
    global memprof
    global shared_raster
    memprof = _MemoryProfiler()
    if r is not None:
        # print("Initializing shared raster in worker")
        shared_raster = r


def _align_partition(
    grid_partition_gdf: gpd.GeoDataFrame,
    raster: xr.Dataset | Callable[[], xr.Dataset] | None,
    aggregations: _Aggregations | str,  # str -> Interpolation method
    pxbuffer: int,
):
    # ? This function is expected to run inside a worker process
    # It heavily utilizes different techniques to reduce memory usage such as
    # Lazy operations, reading only necessary data, and cleaning up memory after use.
    # Shared in-memory raster datasets are used when possible to avoid duplicating large datasets in memory.
    # Shared raster datasets only work when using the "fork" start method for multiprocessing.

    # Strategy for each cell:
    # 1. Correct the geometry to account for the antimeridian
    # 2. Cop the dataset and load the data into memory
    # 3. Aggregate the data
    # 4. Store the aggregated data in the correct location in the output arrays

    # Steps 1-3 are done in different Processes, since all of these steps are CPU bound
    # Even the reading part is CPU bound, since it involves applying masks and cropping the data
    # Also, the backend (icechunk or zarr) can do some computations, like combining chunks together

    # For the Arcticdem on Hex3 grid the avg. time per task is:
    # 1. ~0.01s
    # 2. ~1-3s
    # 3. ~10-20s
    # And for Hex6:
    # 1. unmeasurable
    # 2. ~0.1-0.3s
    # 3. ~0.05s
    # => There is a shift towards step 2 being the bottleneck for higher resolution grids, thus a simple loop becomes
    # faster than a processpoolexecutor

    assert memprof is not None, "Memory profiler is not initialized in worker"
    memprof.log_memory("Before reading partial raster", log=False)

    need_to_close_raster = False

    if raster is None:
        assert shared_raster is not None, "Shared raster is not initialized in worker"
        # print("Using shared raster in worker")
        raster = shared_raster
    elif callable(raster):
        # print("Loading raster in partition")
        raster = raster()
        need_to_close_raster = True
    # else:
    # print("Using provided raster in partition")

    # Partition the raster if necessary
    is_raster_in_memory = isinstance(raster, xr.Dataset) and all(
        isinstance(raster[var].data, np.ndarray) for var in raster.data_vars
    )
    if is_raster_in_memory:
        partial_raster = raster
    else:
        partial_extent = odc.geo.BoundingBox(*grid_partition_gdf.total_bounds, crs=grid_partition_gdf.crs)
        partial_extent = partial_extent.buffered(
            raster.odc.geobox.resolution.x * pxbuffer,
            raster.odc.geobox.resolution.y * pxbuffer,
        )  # buffer by pxbuffer pixels
        with stopwatch("Cropping raster to partition extent", log=False):
            try:
                partial_raster: xr.Dataset = raster.odc.crop(partial_extent, apply_mask=False).compute()
            except Exception as e:
                print(f"Error cropping raster to partition extent: {e}")
                raise e

        if partial_raster.nbytes / 1e9 > 20:
            print(
                f"{os.getpid()}: WARNING! Partial raster size is larger than 20GB:"
                f" {partial_raster.nbytes / 1e9:.2f} GB ({len(grid_partition_gdf)} cells)."
                f" This may lead to out-of-memory errors."
            )
    memprof.log_memory("After reading partial raster", log=False)

    if isinstance(aggregations, str):
        cell_ids = grids.convert_cell_ids(grid_partition_gdf)
        if grid_partition_gdf.crs.to_epsg() == 4326:
            centroids = grid_partition_gdf.geometry.apply(antimeridian.fix_shape).apply(antimeridian.centroid)
            cx = centroids.apply(lambda p: p.x)
            cy = centroids.apply(lambda p: p.y)
        else:
            centroids = grid_partition_gdf.geometry.centroid
            cx = centroids.x
            cy = centroids.y
        interp_x = xr.DataArray(cx, dims=["cell_ids"], coords={"cell_ids": cell_ids})
        interp_y = xr.DataArray(cy, dims=["cell_ids"], coords={"cell_ids": cell_ids})
        ydim = "latitude" if "latitude" in raster.dims else "y"
        xdim = "longitude" if "longitude" in raster.dims else "x"
        interp_coords = {ydim: interp_y, xdim: interp_x}

        # ?: Cubic does not work with NaNs in xarray interp
        with stopwatch("Interpolating data to grid centroids", log=False):
            # Fill the nan
            y_is_rev = partial_raster.indexes[ydim].is_monotonic_decreasing
            if y_is_rev:
                partial_raster = partial_raster.sortby(ydim)
            partial_raster = partial_raster.interpolate_na(
                dim=ydim,
                method=aggregations,
            ).interpolate_na(
                dim=xdim,
                method=aggregations,
            )
            if y_is_rev:
                partial_raster = partial_raster.sortby(ydim)
            ongrid = partial_raster.interp(interp_coords, method=aggregations, kwargs={"fill_value": np.nan})
        memprof.log_memory("After interpolating data", log=False)
    else:
        others_shape = tuple(
            [raster.sizes[dim] for dim in raster.dims if dim not in ["y", "x", "latitude", "longitude"]]
        )
        ongrid_shape = (
            len(grid_partition_gdf),
            len(raster.data_vars),
            len(aggregations),
            *others_shape,
        )
        ongrid = np.full(ongrid_shape, np.nan, dtype=np.float32)

        for i, (idx, row) in enumerate(grid_partition_gdf.iterrows()):
            try:
                cell_data = _process_geom(row.geometry, partial_raster, aggregations)
            except (SystemError, SystemExit, KeyboardInterrupt) as e:
                raise e
            except Exception as e:
                print(f"Error processing cell {row['cell_id']}: {e}")
                continue
            ongrid[i, ...] = cell_data

        cell_ids = grids.convert_cell_ids(grid_partition_gdf)
        dims = ["cell_ids", "variables", "aggregations"]
        coords = {
            "cell_ids": cell_ids,
            "variables": list(raster.data_vars),
            "aggregations": aggregations.aggnames(),
        }
        for dim in set(raster.dims) - {"y", "x", "latitude", "longitude"}:
            dims.append(dim)
            coords[dim] = raster.coords[dim]

        ongrid = xr.DataArray(ongrid, dims=dims, coords=coords).to_dataset("variables")

    partial_raster.close()
    del partial_raster
    if need_to_close_raster:
        raster.close()
        del raster
    gc.collect()
    memprof.log_memory("After cleaning", log=False)

    # print("Finished processing partition")
    # print("### Stopwatch summary ###\n")
    # print(stopwatch.summary())
    # print("### Memory summary ###\n")
    # print(memprof.summary())
    # print("#########################")

    return ongrid


@stopwatch("Aligning data with grid")
def _align_data(
    grid_gdf: gpd.GeoDataFrame | list[gpd.GeoDataFrame],
    raster: xr.Dataset | Callable[[], xr.Dataset],
    aggregations: _Aggregations | str,
    n_partitions: int | None,
    concurrent_partitions: int,
    pxbuffer: int,
):
    # ? Logic for memory management of raster dataset in multiprocessing scenarios
    # There are multiple possible scenarios for the raster dataset:
    # Raster can be either:
    # 1. A in-memory xarray.Dataset object
    # 2. A callable that returns an xarray.Dataset object
    # 3. A lazy xarray.Dataset object (e.g., backed by zarr or icechunk)
    # Concurrent partitions can be either:
    # A. 1 (no multiprocessing)
    # B. >1 (multiprocessing with multiple workers in fork)
    # C. >1 (multiprocessing with multiple workers in spawn/forkserver)

    # This results into potential working modes:
    # 1. Work on the complete raster in memory (no partitioning of the raster is needed)
    #    A. Pass it to the function
    #    B. Make a raster handle in each worker
    #    C. Move it into a shared memory buffer (NOT Implemented)
    # 2. & 3. Work on partitions of the raster
    #    A. Load a partition into memory in each task
    #    B. & C. Make a raster handle in each worker and load partitions into memory in each task
    # ! Note, this is not YET implemented, but this is how it SHOULD be implemented

    partial_ongrids = []

    if isinstance(grid_gdf, list):
        n_partitions = len(grid_gdf)
        grid_partitions = grid_gdf
    else:
        assert n_partitions is not None, "n_partitions must be provided when grid_gdf is not a list"
        grid_partitions = partition_grid(grid_gdf, n_partitions)

    if n_partitions < concurrent_partitions:
        print(f"Adjusting concurrent_partitions from {concurrent_partitions} to {n_partitions}")
        concurrent_partitions = n_partitions

    if concurrent_partitions <= 1:
        _init_worker(None)  # No need to use a shared raster, since the processing is done in the main process
        for i, grid_partition in enumerate(grid_partitions):
            # print(f"Processing partition {i + 1}/{n_partitions} with {len(grid_partition)} cells")
            part_ongrid = _align_partition(
                grid_partition,
                raster,
                aggregations,
                pxbuffer,
            )
            partial_ongrids.append(part_ongrid)
    else:
        is_raster_in_memory = isinstance(raster, xr.Dataset) and all(
            isinstance(raster[var].data, np.ndarray) for var in raster.data_vars
        )
        is_mpfork = mp.get_start_method(allow_none=True) == "fork"
        taskarg = None if is_raster_in_memory else raster
        # For spawn or forkserver, we need to copy the raster into each worker
        workerargs = (None if not is_raster_in_memory or not is_mpfork else raster,)
        # For mp start method fork, we can share the raster dataset between workers
        if is_mpfork and is_raster_in_memory:
            assert isinstance(raster, xr.Dataset)  # satisfy type checker, but this is already checked above
            _init_worker(raster)

        with ProcessPoolExecutor(
            max_workers=concurrent_partitions,
            initializer=_init_worker,
            initargs=workerargs,
        ) as executor:
            futures = {}
            for i, grid_partition in enumerate(grid_partitions):
                futures[
                    executor.submit(
                        _align_partition,
                        grid_partition,
                        taskarg,
                        aggregations,
                        pxbuffer,
                    )
                ] = i

            for future in track(
                as_completed(futures),
                total=len(futures),
                description="Processing grid partitions...",
            ):
                try:
                    i = futures[future]
                    # print(f"Processed partition {i + 1}/{len(futures)}")
                    part_ongrid = future.result()
                    partial_ongrids.append(part_ongrid)
                except Exception as e:
                    print(f"Error processing partition {i}: {e}")
                    raise e

    ongrid = xr.concat(partial_ongrids, dim="cell_ids")
    return ongrid


def aggregate_raster_into_grid(
    raster: xr.Dataset | Callable[[], xr.Dataset],
    grid_gdf: gpd.GeoDataFrame | list[gpd.GeoDataFrame],
    aggregations: _Aggregations | Literal["nearest", "linear", "cubic", "interpolate"],
    grid: Grid,
    level: int,
    n_partitions: int | None = 20,
    concurrent_partitions: int = 5,
    pxbuffer: int = 15,
):
    """Aggregate raster data into grid cells.

    Args:
        raster (xr.Dataset | Callable[[], xr.Dataset]): Raster data or a function that returns it.
        grid_gdf (gpd.GeoDataFrame | list[gpd.GeoDataFrame]): The grid to aggregate into.
            If a list of GeoDataFrames is provided, each will be processed as a separate partition.
            No further partitioning will be done and the n_partitions argument will be ignored.
        aggregations (_Aggregations | Literal[["nearest", "linear", "cubic", "interpolate"]):
            The aggregations to perform.
            If a string is provided, interpolation will be used with the specified method.
            If "interpolate" is provided, the nearest neighbor interpolation will be used.
            Supported methods are "nearest", "linear", and "cubic".
        grid (Grid): The type of grid to use.
        level (int): The level of the grid.
        n_partitions (int | None, optional): Number of partitions to divide the grid into. Defaults to 20.
        concurrent_partitions (int, optional): Maximum number of worker processes when processing partitions.
            Defaults to 5.
        pxbuffer (int, optional): Pixel buffer around each grid cell. Defaults to 15.

    Returns:
        xr.Dataset: Aggregated data aligned with the grid.

    """
    ongrid = _align_data(
        grid_gdf,
        raster,
        aggregations if aggregations != "interpolate" else "nearest",
        n_partitions=n_partitions,
        concurrent_partitions=concurrent_partitions,
        pxbuffer=pxbuffer,
    )
    # Dims of aligned: (cell_ids, variables, aggregations, other dims...)

    if callable(raster) and not isinstance(raster, xr.Dataset):
        raster = raster()

    gridinfo = {
        "grid_name": "h3" if grid == "hex" else grid,
        "level": level,
    }
    if grid == "healpix":
        gridinfo["indexing_scheme"] = "nested"
    ongrid.cell_ids.attrs = gridinfo
    for var in raster.data_vars:
        ongrid[var].attrs = raster[var].attrs

    ongrid = xdggs.decode(ongrid)
    return ongrid
