balanced_allocation.py

"""Balanced service area partitioning.

Partitions a cost surface into territories of roughly equal cost-weighted
area.  Each cell is assigned to the source whose biased cost-distance is
lowest, where biases are iteratively adjusted so that no source's territory
is disproportionately large or small.

Algorithm
---------
1. Run ``cost_distance`` once per source to get N cost surfaces.
2. Assign each cell to ``argmin(cost[i] + bias[i])``.
3. Compute each territory's cost-weighted area (sum of friction values).
4. Adjust biases proportionally: increase for over-large territories,
   decrease for under-served ones.
5. Repeat 2-4 until convergence or ``max_iterations``.
"""

from __future__ import annotations

import numpy as np
import xarray as xr

from xrspatial.cost_distance import cost_distance
from xrspatial.utils import _validate_raster

try:
    import dask.array as da
except ImportError:
    da = None

try:
    import cupy
except ImportError:
    class cupy:  # type: ignore[no-redef]
        ndarray = False


def _to_numpy(arr):
    """Convert any array (numpy, cupy, dask) to a numpy array."""
    if da is not None and isinstance(arr, da.Array):
        arr = arr.compute()
    if hasattr(arr, 'get'):  # cupy
        return arr.get()
    return np.asarray(arr)


def _as_numpy(arr):
    """Convert a computed array (numpy or cupy) to numpy."""
    if hasattr(arr, 'get'):
        return arr.get()
    return np.asarray(arr)


def _extract_sources(raster, target_values):
    """Return sorted array of unique source IDs from the raster."""
    data = _to_numpy(raster.data)
    if len(target_values) > 0:
        ids = np.asarray(target_values, dtype=np.float64)
    else:
        mask = np.isfinite(data) & (data != 0)
        ids = np.unique(data[mask])
    return ids[np.isfinite(ids)]


def _make_single_source_raster(raster, source_id):
    """Create a raster with only one source value, rest zero."""
    data = raster.data
    # Build mask for this source
    if da is not None and isinstance(data, da.Array):
        single = da.where(data == source_id, source_id, 0.0)
    elif hasattr(data, 'get'):  # cupy
        import cupy as cp
        single = cp.where(data == source_id, source_id, 0.0)
    else:
        single = np.where(data == source_id, source_id, 0.0)

    return xr.DataArray(
        single,
        coords=raster.coords,
        dims=raster.dims,
        attrs=raster.attrs,
    )


def _sum_where(friction_data, alloc_data, source_id):
    """Sum friction values where allocation matches source_id.

    Returns a Python float.  Works with numpy, cupy, and dask arrays.
    """
    if da is not None and isinstance(friction_data, da.Array):
        masked = da.where(alloc_data == source_id, friction_data, 0.0)
        return float(masked.sum().compute())
    elif hasattr(friction_data, 'get'):  # cupy
        import cupy as cp
        masked = cp.where(alloc_data == source_id, friction_data, 0.0)
        return float(cp.asnumpy(masked.sum()))
    else:
        masked = np.where(alloc_data == source_id, friction_data, 0.0)
        return float(masked.sum())


def _allocate_from_costs(cost_stack, source_ids, fill_value=np.nan):
    """Assign each cell to the source with lowest cost.

    Parameters
    ----------
    cost_stack : list of arrays
        Per-source cost-distance arrays (same shape).
    source_ids : 1-D numpy array
        Source ID for each layer in cost_stack.
    fill_value : float
        Value for cells unreachable from any source.

    Returns
    -------
    allocation : array
        2-D array of source IDs.
    """
    # Stack along axis 0 -> shape (N, H, W)
    first = cost_stack[0]
    if da is not None and isinstance(first, da.Array):
        stacked = da.stack(cost_stack, axis=0)
        # Replace NaN with inf for argmin
        stacked_clean = da.where(da.isnan(stacked), np.inf, stacked)
        best_idx = da.argmin(stacked_clean, axis=0).compute()
        best_idx = _as_numpy(best_idx)
    elif hasattr(first, 'get'):  # cupy
        import cupy as cp
        stacked = cp.stack(cost_stack, axis=0)
        stacked_clean = cp.where(cp.isnan(stacked), cp.inf, stacked)
        best_idx = cp.asnumpy(cp.argmin(stacked_clean, axis=0))
    else:
        stacked = np.stack(cost_stack, axis=0)
        stacked_clean = np.where(np.isnan(stacked), np.inf, stacked)
        best_idx = np.argmin(stacked_clean, axis=0)

    # Map index back to source ID
    alloc = source_ids[best_idx]

    # Mark cells that are unreachable from all sources
    if da is not None and isinstance(first, da.Array):
        all_nan = da.all(da.isnan(da.stack(cost_stack, axis=0)), axis=0)
        all_nan = _as_numpy(all_nan.compute())
    elif hasattr(first, 'get'):
        import cupy as cp
        all_nan = cp.asnumpy(
            cp.all(cp.isnan(cp.stack(cost_stack, axis=0)), axis=0)
        )
    else:
        all_nan = np.all(np.isnan(np.stack(cost_stack, axis=0)), axis=0)

    alloc = alloc.astype(np.float32)
    alloc[all_nan] = fill_value

    return alloc


def _allocate_biased(cost_stack, biases, source_ids, fill_value=np.nan):
    """Assign each cell to source with lowest (cost + bias).

    Like _allocate_from_costs but adds per-source bias before argmin.
    """
    first = cost_stack[0]
    n = len(cost_stack)

    if da is not None and isinstance(first, da.Array):
        layers = []
        for i in range(n):
            layer = da.where(da.isnan(cost_stack[i]), np.inf,
                             cost_stack[i] + biases[i])
            layers.append(layer)
        stacked = da.stack(layers, axis=0)
        best_idx = da.argmin(stacked, axis=0).compute()
        best_idx = _as_numpy(best_idx)
    elif hasattr(first, 'get'):
        import cupy as cp
        layers = []
        for i in range(n):
            layer = cp.where(cp.isnan(cost_stack[i]), cp.inf,
                             cost_stack[i] + biases[i])
            layers.append(layer)
        stacked = cp.stack(layers, axis=0)
        best_idx = cp.asnumpy(cp.argmin(stacked, axis=0))
    else:
        layers = []
        for i in range(n):
            layer = np.where(np.isnan(cost_stack[i]), np.inf,
                             cost_stack[i] + biases[i])
            layers.append(layer)
        stacked = np.stack(layers, axis=0)
        best_idx = np.argmin(stacked, axis=0)

    alloc = source_ids[best_idx].astype(np.float32)

    # Mark unreachable cells
    if da is not None and isinstance(first, da.Array):
        all_nan = da.all(da.isnan(da.stack(cost_stack, axis=0)), axis=0)
        all_nan = _as_numpy(all_nan.compute())
    elif hasattr(first, 'get'):
        import cupy as cp
        all_nan = cp.asnumpy(
            cp.all(cp.isnan(cp.stack(cost_stack, axis=0)), axis=0)
        )
    else:
        all_nan = np.all(np.isnan(np.stack(cost_stack, axis=0)), axis=0)

    alloc[all_nan] = fill_value
    return alloc


def balanced_allocation(
    raster: xr.DataArray,
    friction: xr.DataArray,
    x: str = "x",
    y: str = "y",
    target_values: list = [],
    max_cost: float = np.inf,
    connectivity: int = 8,
    tolerance: float = 0.05,
    max_iterations: int = 100,
    learning_rate: float = 0.5,
) -> xr.DataArray:
    """Partition a cost surface into balanced service territories.

    Assigns each cell to a source such that all territories have roughly
    equal cost-weighted area (sum of friction values).  This extends
    standard cost-distance allocation by iteratively adjusting per-source
    biases until the workload is balanced.

    Parameters
    ----------
    raster : xr.DataArray
        2-D source raster.  Source pixels are identified by non-zero
        finite values (or values in *target_values*).  Each unique value
        is treated as a separate source.
    friction : xr.DataArray
        2-D friction (cost) surface.  Must have the same shape and
        coordinates as *raster*.
    x : str, default='x'
        Name of the x coordinate.
    y : str, default='y'
        Name of the y coordinate.
    target_values : list, optional
        Specific pixel values in *raster* to treat as sources.
        If empty, all non-zero finite pixels are sources.
    max_cost : float, default=np.inf
        Maximum accumulated cost passed to ``cost_distance``.
    connectivity : int, default=8
        Pixel connectivity (4 or 8) passed to ``cost_distance``.
    tolerance : float, default=0.05
        Convergence threshold.  The loop stops when every territory's
        cost-weighted area is within ``tolerance`` of the mean (as a
        fraction of the mean).
    max_iterations : int, default=100
        Maximum number of bias-adjustment iterations.
    learning_rate : float, default=0.5
        Controls how aggressively biases are updated each iteration.
        Smaller values are more stable; larger values converge faster.

    Returns
    -------
    xr.DataArray
        2-D array of source IDs (float32).  Each cell contains the ID
        of the source it is assigned to.  Unreachable cells are NaN.
    """
    _validate_raster(raster, func_name='balanced_allocation', name='raster')
    _validate_raster(friction, func_name='balanced_allocation',
                     name='friction')
    if raster.shape != friction.shape:
        raise ValueError("raster and friction must have the same shape")
    if raster.dims != (y, x):
        raise ValueError(
            f"raster.dims should be ({y!r}, {x!r}), got {raster.dims}"
        )
    if connectivity not in (4, 8):
        raise ValueError("connectivity must be 4 or 8")
    if tolerance <= 0:
        raise ValueError("tolerance must be positive")
    if max_iterations < 1:
        raise ValueError("max_iterations must be >= 1")

    source_ids = _extract_sources(raster, target_values)
    n_sources = len(source_ids)

    if n_sources == 0:
        out = np.full(raster.shape, np.nan, dtype=np.float32)
        return xr.DataArray(out, coords=raster.coords, dims=raster.dims,
                            attrs=raster.attrs)

    if n_sources == 1:
        # Only one source: every reachable cell goes to it
        cd = cost_distance(raster, friction, x=x, y=y,
                           target_values=list(target_values),
                           max_cost=max_cost, connectivity=connectivity)
        cd_np = _to_numpy(cd.data)
        out = np.where(np.isfinite(cd_np), source_ids[0], np.nan)
        return xr.DataArray(out.astype(np.float32), coords=raster.coords,
                            dims=raster.dims, attrs=raster.attrs)

    # Step 1: compute per-source cost-distance surfaces
    cost_surfaces = []  # list of raw data arrays (numpy/cupy/dask)
    for sid in source_ids:
        single = _make_single_source_raster(raster, sid)
        cd = cost_distance(single, friction, x=x, y=y,
                           max_cost=max_cost, connectivity=connectivity)
        cost_surfaces.append(cd.data)

    # Step 2: get friction data for weighting
    fric_data = friction.data
    if da is not None and isinstance(fric_data, da.Array):
        fric_np = fric_data.compute()
        if hasattr(fric_np, 'get'):
            fric_np = fric_np.get()
    elif hasattr(fric_data, 'get'):
        fric_np = fric_data.get()
    else:
        fric_np = np.asarray(fric_data)

    # Replace non-positive and NaN friction with 0 for weighting
    fric_weight = np.where(np.isfinite(fric_np) & (fric_np > 0), fric_np, 0.0)

    # Step 3: iterative balancing
    biases = np.zeros(n_sources, dtype=np.float64)

    for iteration in range(max_iterations):
        # Allocate with current biases
        alloc = _allocate_biased(cost_surfaces, biases, source_ids)

        # Compute per-territory cost-weighted area
        weights = np.array([
            float(np.sum(fric_weight[alloc == sid]))
            for sid in source_ids
        ])

        # Handle sources with no reachable cells
        total = weights.sum()
        if total == 0:
            break

        target_weight = total / n_sources

        # Check convergence: max relative deviation
        nonzero = weights > 0
        if np.any(nonzero):
            max_dev = np.max(np.abs(weights - target_weight)) / target_weight
            if max_dev <= tolerance:
                break

        # Update biases proportionally
        for i in range(n_sources):
            if target_weight > 0:
                deviation = (weights[i] - target_weight) / target_weight
                biases[i] += learning_rate * deviation

    # Build output DataArray
    return xr.DataArray(
        alloc.astype(np.float32),
        coords=raster.coords,
        dims=raster.dims,
        attrs=raster.attrs,
    )