"""Volume change calculation tools (aimed for glaciers)."""
from __future__ import annotations

import warnings
from typing import Any, Callable

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio.fill
import scipy.interpolate
from geoutils.raster import (
    RasterType,
    get_array_and_mask,
    get_mask_from_array,
    get_valid_extent,
)
from tqdm import tqdm

try:
    import cv2

    _has_cv2 = True
except ImportError:
    _has_cv2 = False

import xdem
from xdem._typing import MArrayf, NDArrayf


def hypsometric_binning(
    ddem: NDArrayf,
    ref_dem: NDArrayf,
    bins: float | np.ndarray[Any, np.dtype[np.floating[Any] | np.integer[Any]]] = 50.0,
    kind: str = "fixed",
    aggregation_function: Callable[[NDArrayf], float] = np.median,
) -> pd.DataFrame:
    """
    Separate the dDEM in discrete elevation bins.
    The elevation bins will be calculated based on all ref_dem valid pixels.
    ddem may contain NaN/masked values over the same area, they will be excluded before the aggregation.

    It is assumed that the dDEM is calculated as 'ref_dem - dem' (not 'dem - ref_dem').

    :param ddem: The dDEM as a 2D or 1D array.
    :param ref_dem: The reference DEM as a 2D or 1D array.
    :param bins: The bin size, count, or array, depending on the binning method ('kind').
    :param kind: The kind of binning to do. Choices: ['fixed', 'count', 'quantile', 'custom'].
    :param aggregation_function: The function to aggregate the elevation values within a bin. Defaults to the median.

    :returns: A Pandas DataFrame with elevation bins and dDEM statistics.
    """
    assert ddem.shape == ref_dem.shape

    # Convert ddem mask into NaN
    ddem, _ = get_array_and_mask(ddem)

    # Extract only the valid values, i.e. valid in ref_dem
    valid_mask = ~get_mask_from_array(ref_dem)
    ddem = np.array(ddem[valid_mask])
    ref_dem = np.array(ref_dem.squeeze()[valid_mask])

    if isinstance(bins, np.ndarray):
        zbins = bins
    elif kind == "fixed":
        zbins = np.arange(ref_dem.min(), ref_dem.max() + bins + 1e-6, step=bins)  # +1e-6 in case min=max (1 point)
    elif kind == "count":
        # Make bins between mean_dem.min() and a little bit above mean_dem.max().
        # The bin count has to be bins + 1 because zbins[0] will be a "below min value" bin, which will be irrelevant.
        zbins = np.linspace(ref_dem.min(), ref_dem.max() + 1e-6 / bins, num=int(bins + 1))
    elif kind == "quantile":
        # Make the percentile steps. The bins + 1 is explained above.
        steps = np.linspace(0, 100, num=int(bins) + 1)
        zbins = np.fromiter((np.percentile(ref_dem, step) for step in steps), dtype=float)
        # The uppermost bin needs to be a tiny amount larger than the highest value to include it.
        zbins[-1] += 1e-6
    elif kind == "custom":
        zbins = bins  # type: ignore
    else:
        raise ValueError(f"Invalid bin kind: {kind}. Choices: ['fixed', 'count', 'quantile', 'custom']")

    # Generate bins and get bin indices from the mean DEM
    indices = np.digitize(ref_dem, bins=zbins)

    # Calculate statistics for each bin.
    # If no values exist, all stats should be nans (except count with should be 0)
    # medians, means, stds, nmads = (np.zeros(shape=bins.shape[0] - 1, dtype=ddem.dtype) * np.nan, ) * 4
    values = np.full(shape=zbins.shape[0] - 1, fill_value=np.nan, dtype=ddem.dtype)
    counts = np.zeros_like(values, dtype=int)
    for i in np.arange(indices.min(), indices.max() + 1):
        values_in_bin = ddem[indices == i]

        # Remove possible Nans
        values_in_bin = values_in_bin[np.isfinite(values_in_bin)]

        # Skip if no values are in the bin.
        if values_in_bin.shape[0] == 0:
            continue

        try:
            values[i - 1] = aggregation_function(values_in_bin)
            counts[i - 1] = values_in_bin.shape[0]
        except IndexError as exception:
            # If custom bins were added, i may exceed the bin range, which will be silently ignored.
            if kind == "custom" and "out of bounds" in str(exception):
                continue
            raise exception

    # Collect the results in a dataframe
    output = pd.DataFrame(
        index=pd.IntervalIndex.from_breaks(zbins), data=np.vstack([values, counts]).T, columns=["value", "count"]
    )

    return output


def interpolate_hypsometric_bins(
    hypsometric_bins: pd.DataFrame,
    value_column: str = "value",
    method: str = "polynomial",
    order: int = 3,
    count_threshold: int | None = None,
) -> pd.DataFrame:
    """
    Interpolate hypsometric bins using any valid Pandas interpolation technique.

    NOTE: It will not extrapolate!

    :param hypsometric_bins: Bins where nans will be interpolated.
    :param value_column: The name of the column in 'hypsometric_bins' to use as values.
    :param method: Any valid Pandas interpolation technique (e.g. 'linear', 'polynomial', 'ffill', 'bfill').
    :param order: The polynomial order to use. Only used if method='polynomial'.
    :param count_threshold: Optional. A pixel count threshold to exclude during the curve fit (requires a 'count'
        column).
    :returns: Bins interpolated with the chosen interpolation method.
    """
    # Copy the bins that will be filled.
    bins = hypsometric_bins.copy()
    # Temporarily set the index to be the midpoint (for interpolation)
    bins.index = bins.index.mid

    # Set all bins that are under a (potentially) specified count to NaN (they should be excluded from interpolation)
    if count_threshold is not None:
        assert "count" in hypsometric_bins.columns, "'count' not a column in the dataframe"
        bins_under_threshold = bins["count"] < count_threshold
        bins.loc[bins_under_threshold, value_column] = np.nan

    # Count number of valid (finite) values
    nvalids = np.count_nonzero(np.isfinite(bins[value_column]))

    if nvalids <= order + 1:
        # Cannot interpolate -> leave as it is
        warnings.warn("Not enough valid bins for interpolation -> returning copy", UserWarning)
        return hypsometric_bins.copy()
    else:
        # Interpolate all bins that are NaN.
        bins[value_column] = bins[value_column].interpolate(method=method, order=order, limit_direction="both")

    # If some points were temporarily set to NaN (to exclude from the interpolation), re-set them.
    if count_threshold is not None:
        bins.loc[bins_under_threshold, value_column] = hypsometric_bins.loc[bins_under_threshold.values, value_column]

    # Return the index to intervals instead of the midpoint.
    bins.index = hypsometric_bins.index

    return bins


def fit_hypsometric_bins_poly(
    hypsometric_bins: pd.DataFrame,
    value_column: str = "value",
    degree: int = 3,
    iterations: int = 1,
    count_threshold: int | None = None,
) -> pd.Series:
    """
    Fit a polynomial to the hypsometric bins.

    :param hypsometric_bins: Bins where nans will be interpolated.
    :param value_column: The name of the column in 'hypsometric_bins' to use as values.
    :param degree: The degree of the polynomial to use.
    :param iterations: The number of iterations to run. \
 At each iteration, values with residuals larger than 3 times the residuals' standard deviation are excluded.
    :param count_threshold: Optional. A pixel count threshold to exclude during the curve fit (requires a 'count'
        column).
    :returns: Bins replaced by the polynomial fit.
    """
    bins = hypsometric_bins.copy()
    bins.index = bins.index.mid

    if count_threshold is not None:
        assert "count" in hypsometric_bins.columns, "'count' not a column in the dataframe"
        bins_under_threshold = bins["count"] < count_threshold
        bins.loc[bins_under_threshold, value_column] = np.nan

    # Remove invalid bins
    valids = np.isfinite(np.asarray(bins[value_column]))

    for _k in range(iterations):

        # Fit polynomial
        x = bins.index[valids]
        y = bins[value_column][valids]
        pcoeff = np.polyfit(x, y, deg=degree)

        # Calculate residuals
        interpolated_values = np.polyval(pcoeff, bins.index)
        residuals = interpolated_values - bins[value_column]
        residuals_std = np.nanstd(residuals.values)

        # Filter outliers further than 3 std
        valids_old = np.copy(valids)
        valids[np.abs(residuals.values) > 3 * residuals_std] = False
        if np.array_equal(valids, valids_old):
            break

    # Save as pandas' DataFrame
    output = pd.DataFrame(
        index=hypsometric_bins.index, data=np.vstack([interpolated_values, bins["count"]]).T, columns=["value", "count"]
    )

    return output


def calculate_hypsometry_area(
    ddem_bins: pd.Series | pd.DataFrame,
    ref_dem: NDArrayf,
    pixel_size: float | tuple[float, float],
    timeframe: str = "reference",
) -> pd.Series:
    """
    Calculate the associated representative area of the given dDEM bins.

    By default, the area bins will be representative of the mean timing between the reference and nonreference DEM:
    elevations = ref_dem - (h_vs_dh_funcion(ref_dem) / 2)
    This can be changed to either "reference":
    elevations = ref_dem
    Or "nonreference":
    elevations = ref_dem - h_vs_dh_function(ref_dem)

    :param ddem_bins: A Series or DataFrame of dDEM values. If a DataFrame is given, the column 'value' will be used.
    :param ref_dem: The reference DEM. This should not have any NaNs.
    :param pixel_size: The xy or (x, y) size of the reference DEM pixels in georeferenced coordinates.
    :param timeframe: The time at which the area bins are representative. Choices: "reference", "nonreference", "mean"

    :returns: The representative area within the given dDEM bins.
    """
    assert not np.any(np.isnan(ref_dem)), "The given reference DEM has NaNs. No NaNs are allowed to calculate area!"

    if timeframe not in ["reference", "nonreference", "mean"]:
        raise ValueError(f"Argument 'timeframe={timeframe}' is invalid. Choices: ['reference', 'nonreference', 'mean']")

    if isinstance(ddem_bins, pd.DataFrame):
        ddem_bins = ddem_bins["value"]

    # For timeframe "mean" or "nonreference", check that ddem_bins values can be interpolated at any altitude
    if timeframe in ["mean", "nonreference"]:
        assert not np.any(np.isnan(ddem_bins.values)), "The dDEM bins cannot contain NaNs. Remove or fill them first."

        # Generate a continuous elevation vs. dDEM function
        ddem_func = scipy.interpolate.interp1d(
            ddem_bins.index.mid, ddem_bins.values, kind="linear", fill_value="extrapolate"
        )

    # Generate average elevations by subtracting half of the dDEM's values to the reference DEM
    if timeframe == "mean":
        elevations = ref_dem - (ddem_func(ref_dem.data) / 2)
    elif timeframe == "reference":
        elevations = ref_dem
    elif timeframe == "nonreference":
        elevations = ref_dem - ddem_func(ref_dem.data)

    # Extract the bins from the dDEM series and compute the frequency of points in the bins.
    bins = np.r_[[ddem_bins.index.left[0]], ddem_bins.index.right]
    bin_counts = np.histogram(elevations, bins=bins)[0]

    # Multiply the bin counts with the pixel area to get the full area.
    bin_area = bin_counts * (pixel_size**2 if not isinstance(pixel_size, tuple) else pixel_size[0] * pixel_size[1])

    # Put this in a series which will be returned.
    output = pd.Series(index=ddem_bins.index, data=bin_area)

    return output


def linear_interpolation(
    array: NDArrayf | MArrayf,
    max_search_distance: int = 10,
    extrapolate: bool = False,
    force_fill: bool = False,
) -> NDArrayf:
    """
    Interpolate a 2D array using rasterio's fillnodata.

    :param array: An array with NaNs or a masked array to interpolate.
    :param max_search_distance: The maximum number of pixels to search in all directions to find values \
to interpolate from. The default is 10.
    :param extrapolate: if False, will remove pixels that have been extrapolated by fillnodata. Default is False.
    :param force_fill: if True, will replace all remaining gaps by the median of all valid values. Default is False.

    :returns: A filled array with no NaNs
    """
    if not _has_cv2:
        raise ValueError("Optional dependency needed. Install 'opencv'")

    # Create a mask for where nans exist
    nan_mask = get_mask_from_array(array)

    interpolated_array = rasterio.fill.fillnodata(
        array.copy(), mask=(~nan_mask).astype("uint8"), max_search_distance=max_search_distance
    )

    # Remove extrapolated values: gaps up to the size of max_search_distance are kept,
    # but surfaces that artificially grow on the edges are removed
    if not extrapolate:
        interp_mask = cv2.morphologyEx(
            (~nan_mask).squeeze().astype("uint8"), cv2.MORPH_CLOSE, kernel=np.ones((max_search_distance - 1,) * 2)
        ).astype("bool")
        if np.ndim(array) == 3:
            interpolated_array[:, ~interp_mask] = np.nan
        else:
            interpolated_array[~interp_mask] = np.nan

    # Fill the nans (values outside of the value boundaries) with the median value
    # This triggers a warning with np.masked_array's because it ignores the mask
    if force_fill:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            interpolated_array[np.isnan(interpolated_array)] = np.nanmedian(array)
    else:
        # If input is masked array, return a masked array
        extrap_mask = interpolated_array != array.data
        if isinstance(array, np.ma.masked_array):
            interpolated_array = np.ma.masked_array(interpolated_array, mask=(nan_mask & ~extrap_mask))

    return interpolated_array.reshape(array.shape)


def hypsometric_interpolation(
    voided_ddem: NDArrayf | MArrayf,
    ref_dem: NDArrayf | MArrayf,
    mask: NDArrayf,
) -> MArrayf:
    """
    Interpolate a dDEM using hypsometric interpolation within the given mask.

    Using `ref_dem`, elevation bins of constant height (hard-coded to 50 m for now) are created.
    Gaps in `voided-ddem`, within the provided `mask`, are filled with the median dDEM value within that bin.

    :param voided_ddem: A dDEM with voids (either an array with nans or a masked array).
    :param ref_dem: The reference DEM in the dDEM comparison.
    :param mask: A mask to delineate the area that will be interpolated (True means hypsometric will be used).
    """
    # Get ddem array with invalid pixels converted to NaN and mask of invalid pixels
    ddem, ddem_mask = get_array_and_mask(voided_ddem)

    # Get ref_dem array with invalid pixels converted to NaN and mask of invalid pixels
    dem, dem_mask = get_array_and_mask(ref_dem)

    # Make sure the mask does not have e.g. the shape (1, height, width)
    mask = mask.squeeze()

    # A mask of inlier values: The union of the mask and the inverted exclusion masks of both rasters.
    inlier_mask = mask & (~ddem_mask & ~dem_mask)
    if np.count_nonzero(inlier_mask) == 0:
        warnings.warn("No valid data found within mask, returning copy", UserWarning)
        return np.ma.masked_array(data=voided_ddem)

    # Estimate the elevation dependent gradient.
    gradient = xdem.volume.hypsometric_binning(ddem[inlier_mask], dem[inlier_mask])

    # Interpolate possible missing elevation bins in 1D - no extrapolation done here
    interpolated_gradient = xdem.volume.interpolate_hypsometric_bins(gradient)

    gradient_model = scipy.interpolate.interp1d(
        interpolated_gradient.index.mid, interpolated_gradient["value"].values, fill_value="extrapolate"
    )

    # Create an idealized dDEM using the relationship between elevation and dDEM
    idealized_ddem = np.zeros_like(dem)
    idealized_ddem[mask] = gradient_model(dem[mask])

    # Replace ddem gaps with idealized hypsometric ddem, but only within mask
    corrected_ddem = np.where(ddem_mask & mask, idealized_ddem, ddem)

    output = np.ma.masked_array(corrected_ddem, mask=~np.isfinite(corrected_ddem))

    assert output is not None

    return output


def local_hypsometric_interpolation(
    voided_ddem: NDArrayf | MArrayf,
    ref_dem: NDArrayf | MArrayf,
    mask: NDArrayf,
    min_coverage: float = 0.2,
    count_threshold: int | None = 1,
    nodata: float | int = -9999,
    plot: bool = False,
) -> MArrayf:
    """
    Interpolate a dDEM using local hypsometric interpolation.
    The algorithm loops through each features in the vector file.

    The dDEM is assumed to have been created as "voided_ddem = reference_dem - other_dem".

    :param voided_ddem: A dDEM with voids (either an array with nans or a masked array).
    :param ref_dem: The reference DEM in the dDEM comparison.
    :param mask: A raster of same shape as voided_ddem and ref_dem, containing a different non-0 pixel value for \
each geometry on which to loop.
    :param min_coverage: Optional. The minimum coverage fraction to be considered for interpolation.
    :param count_threshold: Optional. A pixel count threshold to exclude during the hypsometric curve fit.
    :param nodata: Optional. No data value to be used for the output masked_array.
    :param plot: Set to True to display intermediate plots.

    :returns: A dDEM with gaps filled by applying a hypsometric interpolation for each geometry in mask, \
for areas filling the min_coverage criterion.
    """
    # Remove any unnecessary dimension
    orig_shape = voided_ddem.shape
    voided_ddem = voided_ddem.squeeze()
    ref_dem = ref_dem.squeeze()
    mask = mask.squeeze()

    # Check that all arrays have same dimensions
    assert voided_ddem.shape == ref_dem.shape == mask.shape

    # Get ddem array with invalid pixels converted to NaN and mask of invalid pixels
    ddem, ddem_mask = get_array_and_mask(voided_ddem)

    # Get ref_dem array with invalid pixels converted to NaN and mask of invalid pixels
    dem, dem_mask = get_array_and_mask(ref_dem)

    # A mask of inlier values: The union of the mask and the inverted exclusion masks of both rasters.
    inlier_mask = (mask != 0) & (~ddem_mask & ~dem_mask)
    if np.count_nonzero(inlier_mask) == 0:
        warnings.warn("No valid data found within mask, returning copy", UserWarning)
        return np.ma.masked_array(voided_ddem)

    if plot:
        plt.matshow(inlier_mask)
        plt.title("inlier mask")
        plt.show()

    # List of indexes to loop on
    geometry_index = np.unique(mask[mask != 0])
    print(f"Found {len(geometry_index):d} geometries")

    # Get fraction of valid pixels for each geometry
    coverage = np.zeros(len(geometry_index))
    for k, index in enumerate(geometry_index):
        local_inlier_mask = inlier_mask & (mask == index)
        total_pixels = np.count_nonzero(mask == index)
        valid_pixels = np.count_nonzero(local_inlier_mask)
        coverage[k] = valid_pixels / float(total_pixels)

    # Filter geometries with too little coverage
    valid_geometry_index = geometry_index[coverage >= min_coverage]
    print(f"Found {len(valid_geometry_index):d} geometries with sufficient coverage")

    idealized_ddem = nodata * np.ones_like(dem)

    for _k, index in enumerate(valid_geometry_index):

        # Mask of valid pixel within geometry
        local_mask = mask == index
        local_inlier_mask = inlier_mask & (local_mask)

        # Estimate the elevation dependent gradient
        gradient = xdem.volume.hypsometric_binning(ddem[local_mask], dem[local_mask])

        # Remove bins with loo low count
        filt_gradient = gradient.copy()
        if count_threshold is not None:
            if count_threshold > 1:
                bins_under_threshold = filt_gradient["count"] < count_threshold
                filt_gradient.loc[bins_under_threshold, "value"] = np.nan

        # Interpolate missing elevation bins
        interpolated_gradient = xdem.volume.interpolate_hypsometric_bins(filt_gradient)

        # At least 2 points needed for interp1d, if not skip feature
        nvalues = len(interpolated_gradient["value"].values)
        if nvalues < 2:
            warnings.warn(
                f"Not enough valid bins for feature with index {index:d} -> skipping interpolation", UserWarning
            )
            continue

        # Create a model for 2D interpolation
        gradient_model = scipy.interpolate.interp1d(
            interpolated_gradient.index.mid, interpolated_gradient["value"].values, fill_value="extrapolate"
        )

        if plot:
            local_ddem = np.where(local_inlier_mask, ddem, np.nan)
            vmax = max(np.abs(np.nanpercentile(local_ddem, [2, 98])))
            rowmin, rowmax, colmin, colmax = get_valid_extent(mask == index)

            plt.figure(figsize=(12, 8))
            plt.subplot(121)
            plt.imshow(
                (mask == index)[rowmin:rowmax, colmin:colmax], cmap="Greys", vmin=0, vmax=2, interpolation="none"
            )

            plt.imshow(
                local_ddem[rowmin:rowmax, colmin:colmax], cmap="RdYlBu", vmin=-vmax, vmax=vmax, interpolation="none"
            )
            plt.colorbar()
            plt.title(f"ddem for geometry # {index:d}")

            plt.subplot(122)
            plt.plot(gradient["value"], gradient.index.mid, label="raw")
            plt.plot(interpolated_gradient, gradient.index.mid, label="interpolated", ls="--")
            plt.xlabel("ddem")
            plt.ylabel("Elevation")
            plt.legend()
            plt.title("Average ddem per elevation bin")
            plt.tight_layout()
            plt.show()

        # Create an idealized dDEM (only considering the dH gradient)
        idealized_ddem[mask == index] = gradient_model(dem[mask == index])

    # Measure the difference between the original dDEM and the idealized dDEM
    assert ddem.shape == idealized_ddem.shape
    ddem_difference = ddem.astype("float32") - idealized_ddem.astype("float32")
    ddem_difference[idealized_ddem == nodata] = np.nan

    # Spatially interpolate the difference between these two products.
    interpolated_ddem_diff = linear_interpolation(np.where(ddem_mask, np.nan, ddem_difference))
    interpolated_ddem_diff[np.isnan(interpolated_ddem_diff)] = 0

    # Correct the idealized dDEM with the difference to the original dDEM.
    corrected_ddem = idealized_ddem + interpolated_ddem_diff

    # Set Nans to nodata
    corrected_ddem[~np.isfinite(corrected_ddem)] = nodata

    output = np.ma.masked_array(
        corrected_ddem, mask=(corrected_ddem == nodata)  # mask=((mask != 0) & (ddem_mask | dem_mask))
    ).reshape(orig_shape)

    assert output is not None

    return output


def get_regional_hypsometric_signal(
    ddem: NDArrayf | MArrayf | RasterType,
    ref_dem: NDArrayf | MArrayf | RasterType,
    glacier_index_map: NDArrayf | RasterType,
    n_bins: int = 20,
    verbose: bool = False,
    min_coverage: float = 0.05,
) -> pd.DataFrame:
    """
    Get the normalized regional hypsometric elevation change signal, read "the general shape of it".

    :param ddem: The dDEM to analyse.
    :param ref_dem: A void-free reference DEM.
    :param glacier_index_map: An array glacier indices of the same shape as the previous inputs.
    :param verbose: Show progress bar.
    n_bins = 20  # TODO: This should be an argument.
    :param n_bins: The number of elevation bins to subdivide each glacier in.

    :returns: A DataFrame of bin statistics, scaled by elevation and elevation change.
    """
    # Extract the array and mask representations of the arrays.
    ddem_arr, ddem_mask = get_array_and_mask(ddem)
    ref_arr, ref_mask = get_array_and_mask(ref_dem)
    glacier_index_map, _ = get_array_and_mask(glacier_index_map)

    # The reference DEM should be void free
    assert np.count_nonzero(ref_mask) == 0, "Reference DEM has voids"

    # The unique indices are the unique glaciers.
    unique_indices = np.unique(glacier_index_map)

    # Create empty (ddem) value and (pixel) count arrays which will be filled iteratively.
    values = np.full((n_bins, unique_indices.shape[0]), fill_value=np.nan, dtype=float)
    counts = np.full((n_bins, unique_indices.shape[0]), fill_value=np.nan, dtype=float)

    # Start a counter of glaciers that are actually processed.
    count = 0
    # Loop over each unique glacier.
    for i in tqdm(np.unique(glacier_index_map), desc="Finding regional signal", disable=(not verbose)):
        # If i ==0, it's assumed to be periglacial.
        if i == 0:
            continue
        # Create a mask representing a particular glacier.
        glacier_values = glacier_index_map == i

        # Stop if the "glacier" is tiny. It might be a cropped glacier outline for example.
        if np.count_nonzero(glacier_values) < 10:
            continue

        # The inlier mask is where that particular glacier is and where nans don't exist.
        inlier_mask = glacier_values & ~ddem_mask

        # Skip if the coverage is below the threshold
        if (np.count_nonzero(inlier_mask) / np.count_nonzero(glacier_values)) < min_coverage:
            continue

        # Extract only the difference and elevation values that correspond to the glacier.
        differences = ddem_arr[inlier_mask]
        elevations = ref_arr[inlier_mask]

        # Run the hypsometric binning.
        try:
            bins = hypsometric_binning(differences, elevations, bins=n_bins, kind="count")
        except ValueError:
            # ValueError: zero-size array to reduction operation minimum which has no identity on "zbins=" call
            continue

        # Min-max scale by elevation.
        bins.index = (bins.index.mid - bins.index.left.min()) / (bins.index.right.max() - bins.index.left.min())

        # Scale by difference.
        bins["value"] = (bins["value"] - np.nanmin(bins["value"])) / (
            np.nanmax(bins["value"]) - np.nanmin(bins["value"])
        )

        # Assign the values and counts to the output array.
        values[:, count] = bins["value"]
        counts[:, count] = bins["count"]

        count += 1

    output = pd.DataFrame(
        data={
            "w_mean": np.nansum(values * counts, axis=1) / np.nansum(counts, axis=1),
            "median": np.nanmedian(values, axis=1),
            "std": np.nanstd(values, axis=1),
            "sigma-1-lower": np.nanpercentile(values, 16, axis=1),
            "sigma-1-upper": np.nanpercentile(values, 84, axis=1),
            "sigma-2-lower": np.nanpercentile(values, 2.5, axis=1),
            "sigma-2-upper": np.nanpercentile(values, 97.5, axis=1),
            "count": np.nansum(counts, axis=1).astype(int),
        },
        index=pd.IntervalIndex.from_breaks(np.linspace(0, 1, n_bins + 1, dtype="float64")),
    )

    return output


def norm_regional_hypsometric_interpolation(
    voided_ddem: NDArrayf | MArrayf | RasterType,
    ref_dem: NDArrayf | MArrayf | RasterType,
    glacier_index_map: NDArrayf | RasterType,
    min_coverage: float = 0.1,
    regional_signal: pd.DataFrame | None = None,
    verbose: bool = False,
    min_elevation_range: float = 0.33,
    idealized_ddem: bool = False,
) -> NDArrayf:
    """
    Interpolate missing values by scaling the normalized regional hypsometric signal to each glacier separately.

    Only missing values are interpolated. The rest of the glacier's values are fixed.

    :param voided_ddem: The voided dDEM to fill NaNs in.
    :param ref_dem: A void-free reference DEM.
    :param glacier_index_map: An array glacier indices of the same shape as the previous inputs.
    :param min_coverage: The minimum fractional coverage of a glacier to interpolate. Defaults to 10%.
    :param regional_signal: A regional signal is already estimate. Otherwise one will be estimated.
    :param verbose: Show progress bars.
    :param min_elevation_range: The minimum allowed min/max bin range to scale a signal from.\
            Default: 1/3 of the elevation range needs to be present.
    :param idealized_ddem: Replace observed glacier values with the hypsometric signal. Good for error assessments.

    :raises AssertionError: If `ref_dem` has voids.

    :returns: A dDEM where glacier's that fit the min_coverage criterion are interpolated.
    """
    # Extract the array and nan parts of the inputs.
    ddem_arr, ddem_nans = get_array_and_mask(voided_ddem)
    ref_arr, ref_nans = get_array_and_mask(ref_dem)
    glacier_index_map, _ = get_array_and_mask(glacier_index_map)

    # The reference DEM should be void free
    assert np.count_nonzero(ref_nans) == 0, "Reference DEM has voids"

    # If the regional signal was not given as an argument, find it from the dDEM.
    if regional_signal is None:
        regional_signal = get_regional_hypsometric_signal(
            ddem=ddem_arr, ref_dem=ref_arr, glacier_index_map=glacier_index_map, verbose=verbose
        )

    # The unique indices are the unique glaciers.
    unique_indices = np.unique(glacier_index_map)

    # Make a copy of the dDEM which will be filled iteratively.
    ddem_filled = ddem_arr.copy()
    # Loop over all glaciers and fill the dDEM accordingly.
    for i in tqdm(unique_indices, desc="Interpolating dDEM", disable=(not verbose)):
        if i == 0:  # i==0 is assumed to mean stable ground.
            continue
        # Create a mask representing a particular glacier.
        glacier_values = glacier_index_map == i

        # The inlier mask is where that particular glacier is and where nans don't exist.
        inlier_mask = glacier_values & ~ddem_nans

        # If the fractional coverage is smaller than the given threshold, skip the glacier.
        if (np.count_nonzero(inlier_mask) / np.count_nonzero(glacier_values)) < min_coverage:
            continue

        # Extract only the finite difference and elevation values that correspond to the glacier.
        differences = ddem_arr[inlier_mask]
        elevations = ref_arr[inlier_mask]

        # Get the reference elevation min and max
        elev_min = ref_arr[glacier_values].min()
        elev_max = ref_arr[glacier_values].max()

        # Copy the signal
        signal = regional_signal["w_mean"].copy()
        # Scale the signal elevation midpoints to the glacier elevation range.
        midpoints = signal.index.mid
        midpoints *= elev_max - elev_min
        midpoints += elev_min
        step = midpoints[1] - midpoints[0]
        # Create an interval structure from the midpoints and the step size.
        signal.index = pd.IntervalIndex.from_arrays(left=midpoints - step / 2, right=midpoints + step / 2)

        # Find the hypsometric bins of the glacier.
        hypsometric_bins = hypsometric_binning(
            ddem=differences,
            ref_dem=elevations,
            bins=np.r_[[signal.index.left[0]], signal.index.right],  # This will generate the same steps as the signal.
            kind="custom",
        )
        bin_stds = hypsometric_binning(
            ddem=differences,
            ref_dem=elevations,
            bins=np.r_[[signal.index.left[0]], signal.index.right],
            kind="custom",
            aggregation_function=np.nanstd,
        )
        # Check which of the bins were non-empty.
        non_empty_bins = np.isfinite(hypsometric_bins["value"])

        non_empty_range = np.sum(non_empty_bins[non_empty_bins].index.length)
        full_range = np.sum(hypsometric_bins.index.length)

        if (non_empty_range / full_range) < min_elevation_range:
            continue

        # A theoretical minimum of 2 bins are needed for the curve fit.
        if np.count_nonzero(non_empty_bins) < 2:
            continue

        # The weights are the squared inverse of the standard deviation of each bin.
        bin_weights = bin_stds["value"].values[non_empty_bins] / np.sqrt(
            hypsometric_bins["count"].values[non_empty_bins]
        )
        bin_weights[bin_weights == 0.0] = 1e-8  # Avoid divide by zero problems.

        # Fit linear coefficients to scale the regional signal to the hypsometric bins properly.
        # The inverse of the pixel counts are used as weights, to properly disregard poorly constrained bins.
        with warnings.catch_warnings():
            # curve_fit will sometimes say "can't estimate covariance". This is okay.
            warnings.filterwarnings("ignore", message="covariance")
            coeffs = scipy.optimize.curve_fit(
                f=lambda x, a, b: a * x + b,  # Estimate a linear function "f(x) = ax + b".
                xdata=signal.values[non_empty_bins],  # The xdata is the normalized regional signal
                ydata=hypsometric_bins["value"].values[non_empty_bins],  # The ydata is the actual values.
                p0=[1, 0],  # The initial guess of a and b (doesn't matter too much)
                sigma=bin_weights,
            )[0]

        # Create a linear model from the elevations and the scaled regional signal.
        model = scipy.interpolate.interp1d(
            signal.index.mid, np.poly1d(coeffs)(signal.values), bounds_error=False, fill_value="extrapolate"
        )

        # Find which values to fill using the model (all nans within the glacier extent)
        if not idealized_ddem:
            values_to_fill = glacier_values & ddem_nans
        # If it should be idealized, replace all glacier values with the model
        else:
            values_to_fill = glacier_values
        # Fill the nans using the scaled regional signal.
        ddem_filled[values_to_fill] = model(ref_arr[values_to_fill])

    return ddem_filled
