"""N-dimensional image manipulation."""

from pathlib import Path

import numpy as np
from scipy.ndimage import label, maximum_filter

###############################################################################
# Operating on connected components
###############################################################################


def largest_connected_component(volume) -> np.ndarray:
    """Return the largest connected component of a 3D array.

    Parameters
    ----------
    volume : numpy.ndarray
        3D boolean array indicating a volume.

    Returns
    -------
    volume : numpy.ndarray
        3D boolean array with only one connected component.

    See Also
    --------
    nilearn.image.largest_connected_component_img : To simply operate the
        same manipulation directly on Nifti images.

    Notes
    -----
    **Handling big-endian in given numpy.ndarray**
    This function changes the existing byte-ordering information to new byte
    order, if the given volume has non-native data type. This operation
    is done inplace to avoid big-endian issues with scipy ndimage module.

    """
    if hasattr(volume, "get_fdata") or isinstance(volume, (str, Path)):
        raise ValueError(
            "Please enter a valid numpy array. For images use "
            "largest_connected_component_img."
        )
    # Get the new byteorder to handle issues like "Big-endian buffer not
    # supported on little-endian compiler" with scipy ndimage label.
    if not volume.dtype.isnative:
        volume.dtype = volume.dtype.newbyteorder("N")

    # We use asarray to be able to work with masked arrays.
    volume = np.asarray(volume)
    labels, label_nb = label(volume)
    if not label_nb:
        raise ValueError("No non-zero values: no connected components")
    if label_nb == 1:
        return volume.astype(bool)
    label_count = np.bincount(labels.ravel().astype(int))
    # discard the 0 label
    label_count[0] = 0
    return labels == label_count.argmax()


def get_border_data(data: np.ndarray, border_size: int) -> np.ndarray:
    """Return the data at the border of an array."""
    return np.concatenate(
        [
            data[:border_size, :, :].ravel(),
            data[-border_size:, :, :].ravel(),
            data[:, :border_size, :].ravel(),
            data[:, -border_size:, :].ravel(),
            data[:, :, :border_size].ravel(),
            data[:, :, -border_size:].ravel(),
        ]
    )


def peak_local_max(
    image: np.ndarray,
    min_distance: int = 10,
    threshold_abs: float = 0,
    threshold_rel: float = 0.1,
    num_peaks=np.inf,
) -> np.ndarray:
    """Find peaks in an image, and return them \
    as coordinates or a boolean array.

    Peaks are the local maxima in a region of `2 * min_distance + 1`
    (i.e. peaks are separated by at least `min_distance`).

    Parameters
    ----------
    image : ndarray of floats
        Input image.

    min_distance : int, default=10
        Minimum number of pixels separating peaks in a region of `2 *
        min_distance + 1` (i.e. peaks are separated by at least
        `min_distance`). To find the maximum number of peaks, use
        `min_distance=1`.

    threshold_abs : float, default=0
        Minimum intensity of peaks.

    threshold_rel : float, default=0.1
        Minimum intensity of peaks calculated as `max(image) * threshold_rel`.

    num_peaks : int, default=np.inf
        Maximum number of peaks. When the number of peaks exceeds `num_peaks`,
        return `num_peaks` peaks based on highest peak intensity.

    Returns
    -------
    output : ndarray or ndarray of bools
        Boolean array shaped like `image`, with peaks represented by True
        values.

    Notes
    -----
    If peaks are flat (i.e. multiple adjacent pixels have identical
    intensities), the coordinates of all such pixels are returned.

    The peak local maximum function returns the coordinates of local peaks
    (maxima) in a image. A maximum filter is used for finding local maxima.
    This operation dilates the original image. After comparison between
    dilated and original image, peak_local_max function returns the
    coordinates of peaks where dilated image = original.

    This code is mostly adapted from scikit image 0.11.3 release.
    Location of file in scikit image: peak_local_max function in
    skimage.feature.peak

    """
    out = np.zeros_like(image, dtype=bool)

    if np.all(image == image.flat[0]):
        return out

    image = image.copy()

    size = 2 * min_distance + 1
    image_max = maximum_filter(image, size=size, mode="constant")

    mask = image == image_max
    image *= mask

    # find top peak candidates above a threshold
    peak_threshold = max(np.max(image.ravel()) * threshold_rel, threshold_abs)

    # get coordinates of peaks
    coordinates = np.argwhere(image > peak_threshold)

    if coordinates.shape[0] > num_peaks:
        intensities = image.flat[
            np.ravel_multi_index(coordinates.transpose(), image.shape)
        ]
        idx_maxsort = np.argsort(intensities)[::-1]
        coordinates = coordinates[idx_maxsort][:num_peaks]

    nd_indices = tuple(coordinates.T)
    out[nd_indices] = True
    return out
