import numpy

from .io import get_persistent_array_mask

try:
    import cupy
    from .cupyutils import log_allocated_gpu_memory
except ImportError:
    CUPY_AVAILABLE = False
    CUPY_MEM_POOL = None
    cupy = numpy
else:
    CUPY_AVAILABLE = True
    CUPY_MEM_POOL = cupy.get_default_memory_pool()
import logging

MASK_PIXELS_AVAILABLE = None

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def get_index_shifted(
    data_shape: tuple,
    center_x: int,
    center_y: int,
):
    """Calculates the x,y index vectors, shifted to recreate the centro-symmetric pixels

    Inputs:
        - data_shape (tuple): shape of the data array (2 dimensional)
        - center_x (int): beam center in the first data dimension
        - center_y (int): beam center in the second data dimension
    Outputs:
        Tuple: pair of shifted x,y index vectors
    """
    y, x = numpy.meshgrid(
        numpy.arange(data_shape[0]),
        numpy.arange(data_shape[1]),
        indexing="ij",
        sparse=True,
    )

    x_ = 2 * int(center_x) - x + 1
    y_ = 2 * int(center_y) - y + 1

    # Think it's correct
    return x_, y_


def get_position_vectors(
    data_shape: tuple,
    use_cupy: bool = False,
):
    """Creates the x,y index vector of a 2-dimensional data array

    Inputs:
        - data_shape (tuple): shape of the data array (2 dimensional)
        - use_cupy (bool): transfer the vectors to the GPU, ready to process using cupy
    Outputs:
        - Tuple: pair of x-y index vector
    """
    y_vector, x_vector = numpy.meshgrid(
        numpy.arange(data_shape[0]),
        numpy.arange(data_shape[1]),
        indexing="ij",
        sparse=True,
    )
    if use_cupy:
        return cupy.asarray(x_vector), cupy.asarray(y_vector)
    return x_vector, y_vector


def shift_position_vectors(
    x_vector: numpy.ndarray,
    y_vector: numpy.ndarray,
    center_x: int,
    center_y: int,
    use_cupy: bool = False,
):
    """Shift a pair of index vectors, to point to centro-symmetric brothers

    Inputs:
        - x_vector (numpy.ndarray) : index vector along the first dimension of array
        - y_vector (numpy.ndarray) : index vector along the second dimension of array
        - center_x (int): beam center in the first data dimension
        - center_y (int): beam center in the second data dimension
        - use_cupy (bool): the return result is transfered to the GPU using cupy
    Outputs:
        - Tuple: pair of shifted x,y index vectors
    """
    center_x = int(center_x)
    center_y = int(center_y)
    x_shifted = 2 * center_x - x_vector + 1
    y_shifted = 2 * center_y - y_vector + 1
    if use_cupy:
        return cupy.asarray(x_shifted), cupy.asarray(y_shifted)
    return x_shifted, y_shifted


def get_mask_limits(
    data_shape: tuple = None,
    center_x: int = None,
    center_y: int = None,
    x_shifted: numpy.ndarray = None,
    y_shifted: numpy.ndarray = None,
    only_x_dim: bool = False,
    only_y_dim: bool = False,
    use_cupy: bool = False,
):
    """
    Mask limits is the primitive mask, it defines the area where the cave can be applied.
    It's the area of the pixels whose centro-symmetric brother falls in the detector surface.
    If the center of the beam matches with the center of the detector, the whole detector area can be caved.
    If the center falls into the edge of the detector, no pixel can be caved.

    Inputs:
        - data_shape (tuple): shape of the data array (2 dimensional)
        - center_x (int): beam center in the first data dimension
        - center_y (int): beam center in the second data dimension
        - x_shifted (numpy.ndarray) : shifted index vector along the first data dimension
        - y_shifted (numpy.ndarray) : shifted index vector along the second data dimension
        - only_x_dim (bool): if True, returns the mask of the first dimension limits
        - only_y_dim (bool): if True, returns the mask of the second dimension limits
        - use_cupy (bool): the return result is transfered to the GPU using cupy
    Outputs:
        - numpy.ndarray / cupy.ndarray: array to mask the pixels whose centrosymmetric brother falls out of the detector area.
    """
    if x_shifted is None or y_shifted is None:
        x_vector, y_vector = get_position_vectors(
            data_shape=data_shape,
            use_cupy=use_cupy,
        )
        x_shifted, y_shifted = shift_position_vectors(
            x_vector=x_vector,
            y_vector=y_vector,
            center_x=center_x,
            center_y=center_y,
            use_cupy=use_cupy,
        )
    x_max = x_shifted.shape[-1]
    y_max = y_shifted.shape[0]
    if use_cupy:
        mask_limits_x_cupy = cupy.logical_and(x_shifted > 0, x_shifted < x_max)
        if only_x_dim:
            return mask_limits_x_cupy
        mask_limits_y_cupy = cupy.logical_and(y_shifted > 0, y_shifted < y_max)
        if only_y_dim:
            return mask_limits_y_cupy
        mask_limits_cupy = mask_limits_x_cupy & mask_limits_y_cupy  # boolean
        return mask_limits_cupy
    mask_limits_x = numpy.logical_and(x_shifted > 0, x_shifted < x_max)
    if only_x_dim:
        return mask_limits_x
    mask_limits_y = numpy.logical_and(y_shifted > 0, y_shifted < y_max)
    if only_y_dim:
        return mask_limits_y
    mask_limits = mask_limits_x & mask_limits_y  # boolean
    return mask_limits


def _mask_caving(
    data_shape: tuple,
    Center_1: int,
    Center_2: int,
    mask_reference: numpy.ndarray = None,
    vertical_symmetry: bool = True,
    horizontal_symmetry: bool = True,
    use_cupy=False,
) -> numpy.ndarray:
    """Data-free method, generates the mask with the pixels available. Those pixels whose intensity can be substitued by its centrosymmetric, horz-symmetric or vert-symmetric brother.
    It's determined by the position of the beam center, (the closer to the array center, the more pixels available), and the mask_reference, that determines which pixels will not be caved.

    Inputs:
        - data_shape (tuple): shape of the data array (2 dimensional)
        - Center_1 (int): beam center in the first data dimension
        - Center_2 (int): beam center in the second data dimension
        - mask_reference (numpy.ndarray): mask with the pixels whose intensity should not be used, that means that their symmetric brother should not be caved
        - vertical_symmetry (bool): the caving performs a vertical flip around the beam center
        - horizontal_symmetry (bool) : the caving performs a horizontal flip around the beam center
        - use_cupy (bool): returns a cupy array
    Outputs:
        - numpy.ndarray, or cupy.ndarray with the available pixels, taken into account the detector limits and the mask reference
    """
    # Dataset-free method
    # With a reference, this method takes ~20 ms for all operations, < 1ms without reference
    y_vector, x_vector = numpy.meshgrid(
        numpy.arange(data_shape[0]),
        numpy.arange(data_shape[1]),
        indexing="ij",
        sparse=True,
    )
    x_shifted = 2 * int(Center_1) - x_vector + 1
    y_shifted = 2 * int(Center_2) - y_vector + 1
    x_max = x_shifted.shape[-1]
    y_max = y_shifted.shape[0]

    if vertical_symmetry and horizontal_symmetry:
        # Centrosymmetric caving (standard)
        mask_limits = numpy.logical_and(
            x_shifted > 0, x_shifted < x_max
        ) & numpy.logical_and(y_shifted > 0, y_shifted < y_max)
        if mask_reference is None:
            mask_caving = mask_limits
        else:
            mask_caving = mask_limits ^ (
                mask_limits & mask_reference[y_shifted, x_shifted]
            )
    elif horizontal_symmetry:
        # Horizontal flipping around the beam
        mask_limits = numpy.logical_and(x_shifted > 0, x_shifted < x_max)
        if mask_reference is None:
            mask_caving = mask_limits
        else:
            mask_caving = mask_limits ^ (
                mask_limits & mask_reference[y_vector, x_shifted]
            )
    elif vertical_symmetry:
        # Vertical flippinf around the beam
        mask_limits = numpy.logical_and(y_shifted > 0, y_shifted < y_max)
        if mask_reference is None:
            mask_caving = mask_limits
        else:
            mask_caving = mask_limits ^ (
                mask_limits & mask_reference[y_shifted, x_vector]
            )
    else:
        mask_limits = numpy.logical_and(
            x_vector > 0, x_vector < x_max
        ) & numpy.logical_and(y_vector > 0, y_vector < y_max)
        if mask_reference is None:
            mask_caving = mask_limits
        else:
            mask_caving = mask_limits ^ (
                mask_limits & mask_reference[y_vector, x_vector]
            )
    if use_cupy:
        return cupy.asarray(mask_caving).astype(bool)
    return mask_caving.astype(bool)


def _mask_to_cave(
    data: numpy.ndarray, Dummy: int = None, mask_static: numpy.ndarray = None
) -> numpy.ndarray:
    """
    Generates the mask with the pixels that we want to substitute. It's determined by a dummy value and/or by a mask.

    Inputs:
        - data (numpy.ndarray): data signal, 2 or 3 dimensional
        - Dummy (int): pixels in data with dummy value will be replaced
        - mask_static (numpy.ndarray): pixels that will be replaced
    Outputs:
        - numpy.ndarray: mask whose pixels we want to replace
    """
    # ~ 1 ms per frame if only Dummy
    # ~ 2 ms per frame if both dummy and mask_static
    # ~0 if only mask_static
    if Dummy and mask_static is not None:
        mask_to_cave = (data == Dummy) + mask_static
    elif Dummy:
        mask_to_cave = data == Dummy
    elif mask_static is not None:
        mask_to_cave = mask_static
    else:
        raise ValueError(
            "Caving has to be determined by a dummy value and/or a static mask"
        )
    return mask_to_cave


def _process_data_caving(
    data: numpy.ndarray,
    pixels_to_cave: numpy.ndarray,
    pixels_to_substitute: numpy.ndarray,
    x_vector: numpy.ndarray,
    y_vector: numpy.ndarray,
    Dummy: int = None,
    mask_dummy_pixels: bool = True,
    use_cupy: bool = False,
):
    """Replace the pixels in data according to a mask. The replaced data correspond to a specific index-transformation, determined by input vectors.

    Inputs:
        - data (numpy.ndarray): data signal, 2 or 3 dimensional
        - mask_to_cave (numpy.ndarray): pixels that will be masked with Dummy or caved (substituted by a symmetric brother)

        - x_vector (numpy.ndarray): index vector used to shift the data along the first dimension of data array
        - y_vector (numpy.ndarray): index vector used to shift the data along the second dimension of data array
        - Dummy (int): in this method, the dummy value is used to replace the pixels whose replacement falls also into the mask_caving
        - use_cupy (bool): to use cupy, all the inputs must be already in the GPU
    Outputs:
        - numpy.ndarray: data with replaced pixels (caved)
    """
    # Works for both numpy, cupy, 2 or 3 dimensional, for all cavings
    mask_complete = pixels_to_cave & pixels_to_substitute
    if use_cupy:
        data_caved = cupy.copy(data)
        if mask_dummy_pixels and Dummy:
            cupy.copyto(data_caved, Dummy, where=pixels_to_cave)
        if data.ndim == 3:
            cupy.copyto(
                data_caved, data_caved[:, y_vector, x_vector], where=mask_complete
            )
        elif data.ndim == 2:
            cupy.copyto(data_caved, data_caved[y_vector, x_vector], where=mask_complete)
    else:
        data_caved = numpy.copy(data)
        if mask_dummy_pixels and Dummy:
            numpy.copyto(data_caved, Dummy, where=pixels_to_cave)
        if data.ndim == 3:
            numpy.copyto(
                data_caved, data_caved[:, y_vector, x_vector], where=mask_complete
            )
        elif data.ndim == 2:
            numpy.copyto(
                data_caved, data_caved[y_vector, x_vector], where=mask_complete
            )
    return data_caved


def _process_dataset_caving_numpy(
    dataset: numpy.ndarray,
    Center_1: int,
    Center_2: int,
    Dummy: int = None,
    mask_static: numpy.ndarray = None,
    mask_reference: numpy.ndarray = None,
    flip_caving: bool = False,
    flip_horizontally_preference: bool = True,
    **kwargs,
) -> numpy.ndarray:
    """Performs the caving on a 3-dimensional dataset using numpy methods

    Inputs:
        - dataset (numpy.ndarray): data signal, 3-dimensional
        - Center_1 (int): beam center in the first data dimension
        - Center_2 (int): beam center in the second data dimension
        - Dummy (int): pixels in data with dummy value will be replaced
        - mask_static (numpy.ndarray): pixels that will be replaced
        - mask_reference (numpy.ndarray): mask with the pixels whose intensity should not be used, that means that their symmetric brother should not be caved
        - flip_caving (bool): if True, adds horizontal-symmetric and vertical-symmetric caving (by default, centro-symmetric caving is always performed)
        - flip_horizontally_preference (bool): if flip_caving is activated, horizontal caving is done before vertical caving
        - **kwargs
    Outputs:
        - numpy.ndarray: data with replaced pixels (caved) using numpy methods
    """
    if dataset.ndim == 2:
        data_shape = dataset.shape
        dataset = dataset[numpy.newaxis, ...]
    elif dataset.ndim == 3:
        data_shape = dataset.shape[1:]

    y_vector, x_vector = numpy.meshgrid(
        numpy.arange(data_shape[0]),
        numpy.arange(data_shape[1]),
        indexing="ij",
        sparse=True,
    )
    x_shifted = 2 * int(Center_1) - x_vector + 1
    y_shifted = 2 * int(Center_2) - y_vector + 1

    mask_pixels_available_centrosymmetric = kwargs.get(
        "mask_pixels_available_centrosymmetric"
    )
    mask_pixels_available_horizontal = kwargs.get("mask_pixels_available_horizontal")
    mask_pixels_available_vertical = kwargs.get("mask_pixels_available_vertical")
    if mask_pixels_available_centrosymmetric is None:
        mask_pixels_available_centrosymmetric = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=True,
            horizontal_symmetry=True,
        )
    if flip_caving and mask_pixels_available_horizontal is None:
        mask_pixels_available_horizontal = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=False,
            horizontal_symmetry=True,
        )
    if flip_caving and mask_pixels_available_vertical is None:
        mask_pixels_available_vertical = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=True,
            horizontal_symmetry=False,
        )

    # This is the mask with the pixels that will be, either masked with Dummy or caved
    mask_to_cave = _mask_to_cave(data=dataset, Dummy=Dummy, mask_static=mask_static)

    data_caved = _process_data_caving(
        data=dataset,
        Dummy=Dummy,
        mask_dummy_pixels=True,
        pixels_to_cave=mask_to_cave,
        pixels_to_substitute=mask_pixels_available_centrosymmetric,
        x_vector=x_shifted,
        y_vector=y_shifted,
        use_cupy=False,
    )
    if flip_caving and Dummy:
        if flip_horizontally_preference:
            # First horizontal caving
            mask_to_cave = data_caved == Dummy
            data_caved = _process_data_caving(
                data=data_caved,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_horizontal,
                x_vector=x_shifted,
                y_vector=y_vector,
                use_cupy=False,
            )
            # Second vertical caving
            mask_to_cave = data_caved == Dummy
            data_caved = _process_data_caving(
                data=data_caved,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_vertical,
                x_vector=x_vector,
                y_vector=y_shifted,
                use_cupy=False,
            )
        else:
            # First vertical caving
            mask_to_cave = data_caved == Dummy
            data_caved = _process_data_caving(
                data=data_caved,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_vertical,
                x_vector=x_vector,
                y_vector=y_shifted,
                use_cupy=False,
            )
            # Second horizontal caving
            mask_to_cave = data_caved == Dummy
            data_caved = _process_data_caving(
                data=data_caved,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_horizontal,
                x_vector=x_shifted,
                y_vector=y_vector,
                use_cupy=False,
            )
    return data_caved


def _process_dataset_caving_cupy(
    dataset: numpy.ndarray,
    Center_1: int,
    Center_2: int,
    Dummy: int = None,
    mask_static: numpy.ndarray = None,
    mask_reference: numpy.ndarray = None,
    flip_caving: bool = False,
    flip_horizontally_preference: bool = True,
    **kwargs,
) -> numpy.ndarray:
    """Performs the caving on a 3-dimensional dataset using cupy methods (frame by frame)

    Inputs:
        - dataset (numpy.ndarray): data signal, 3-dimensional
        - Center_1 (int): beam center in the first data dimension
        - Center_2 (int): beam center in the second data dimension
        - Dummy (int): pixels in data with dummy value will be replaced
        - mask_static (numpy.ndarray): pixels that will be replaced
        - mask_reference (numpy.ndarray): mask with the pixels whose intensity should not be used, that means that their symmetric brother should not be caved
        - flip_caving (bool): if True, adds horizontal-symmetric and vertical-symmetric caving (by default, centro-symmetric caving is always performed)
        - flip_horizontally_preference (bool): if flip_caving is activated, horizontal caving is done before vertical caving
        - **kwargs
    Outputs:
        - numpy.ndarray: data with replaced pixels (caved) using numpy methods
    """
    log_allocated_gpu_memory()
    if dataset.ndim == 2:
        data_shape = dataset.shape
        dataset = dataset[numpy.newaxis, ...]
    elif dataset.ndim == 3:
        data_shape = dataset.shape[1:]

    y_vector, x_vector = numpy.meshgrid(
        numpy.arange(data_shape[0]),
        numpy.arange(data_shape[1]),
        indexing="ij",
        sparse=True,
    )
    x_shifted = 2 * int(Center_1) - x_vector + 1
    y_shifted = 2 * int(Center_2) - y_vector + 1

    x_shifted_cupy = cupy.asarray(x_shifted)
    y_shifted_cupy = cupy.asarray(y_shifted)
    x_vector_cupy = cupy.asarray(x_vector)
    y_vector_cupy = cupy.asarray(y_vector)

    mask_pixels_available_centrosymmetric = kwargs.get(
        "mask_pixels_available_centrosymmetric"
    )
    mask_pixels_available_horizontal = kwargs.get("mask_pixels_available_horizontal")
    mask_pixels_available_vertical = kwargs.get("mask_pixels_available_vertical")
    if mask_pixels_available_centrosymmetric is None:
        mask_pixels_available_centrosymmetric = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=True,
            horizontal_symmetry=True,
        )
    if flip_caving and mask_pixels_available_horizontal is None:
        mask_pixels_available_horizontal = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=False,
            horizontal_symmetry=True,
        )
    if flip_caving and mask_pixels_available_vertical is None:
        mask_pixels_available_vertical = _mask_caving(
            data_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=True,
            horizontal_symmetry=False,
        )

    if mask_static is not None:
        mask_static_cupy = cupy.asarray(mask_static)
    else:
        mask_static_cupy = None

    mask_pixels_available_centrosymmetric_cupy = cupy.asarray(
        mask_pixels_available_centrosymmetric
    )
    if flip_caving:
        mask_pixels_available_horizontal_cupy = cupy.asarray(
            mask_pixels_available_horizontal
        )
        mask_pixels_available_vertical_cupy = cupy.asarray(
            mask_pixels_available_vertical
        )
    else:
        mask_pixels_available_horizontal_cupy = None
        mask_pixels_available_vertical_cupy = None

    dataset_caved = numpy.zeros_like(dataset)
    for index_frame, data in enumerate(dataset):
        data_cupy = cupy.asarray(data)
        mask_to_cave = _mask_to_cave(data_cupy, Dummy, mask_static_cupy)
        # mask_complete = mask_to_cave & mask_pixels_available_centrosymmetric_cupy
        data_cupy = _process_data_caving(
            data=data_cupy,
            Dummy=Dummy,
            mask_dummy_pixels=True,
            pixels_to_cave=mask_to_cave,
            pixels_to_substitute=mask_pixels_available_centrosymmetric_cupy,
            x_vector=x_shifted_cupy,
            y_vector=y_shifted_cupy,
            use_cupy=True,
        )
        if flip_caving and Dummy:
            if flip_horizontally_preference:
                # First horizontal caving
                mask_to_cave = data_cupy == Dummy
                data_cupy = _process_data_caving(
                    data=data_cupy,
                    Dummy=Dummy,
                    mask_dummy_pixels=False,
                    pixels_to_cave=mask_to_cave,
                    pixels_to_substitute=mask_pixels_available_horizontal_cupy,
                    x_vector=x_shifted_cupy,
                    y_vector=y_vector_cupy,
                    use_cupy=True,
                )
                # Second vertical caving
                mask_to_cave = data_cupy == Dummy
                data_cupy = _process_data_caving(
                    data=data_cupy,
                    Dummy=Dummy,
                    mask_dummy_pixels=False,
                    pixels_to_cave=mask_to_cave,
                    pixels_to_substitute=mask_pixels_available_vertical_cupy,
                    x_vector=x_vector_cupy,
                    y_vector=y_shifted_cupy,
                    use_cupy=True,
                )
            else:
                # First vertical caving
                mask_to_cave = data_cupy == Dummy
                data_cupy = _process_data_caving(
                    data=data_cupy,
                    Dummy=Dummy,
                    mask_dummy_pixels=False,
                    pixels_to_cave=mask_to_cave,
                    pixels_to_substitute=mask_pixels_available_vertical_cupy,
                    x_vector=x_vector_cupy,
                    y_vector=y_shifted_cupy,
                    use_cupy=True,
                )
                # Second horizontal caving
                mask_to_cave = data_cupy == Dummy
                data_cupy = _process_data_caving(
                    data=data_cupy,
                    Dummy=Dummy,
                    mask_dummy_pixels=False,
                    pixels_to_cave=mask_to_cave,
                    pixels_to_substitute=mask_pixels_available_horizontal_cupy,
                    x_vector=x_shifted_cupy,
                    y_vector=y_vector_cupy,
                    use_cupy=True,
                )
        dataset_caved[index_frame] = data_cupy.get()
    return dataset_caved


def _process_data_caving_cupy(
    data_cupy: cupy.ndarray,
    mask_pixels_available_centrosymmetric_cupy: cupy.ndarray,
    x_shifted_cupy: cupy.ndarray,
    y_shifted_cupy: cupy.ndarray,
    Dummy: int = None,
    mask_static_cupy: cupy.ndarray = None,
    flip_caving: bool = False,
    flip_horizontally_preference: bool = True,
    mask_pixels_available_horizontal_cupy: cupy.ndarray = None,
    mask_pixels_available_vertical_cupy: cupy.ndarray = None,
    x_vector_cupy: cupy.ndarray = None,
    y_vector_cupy: cupy.ndarray = None,
) -> cupy.ndarray:
    """Performs the caving on a 3-dimensional dataset using cupy methods (frame by frame)

    Inputs:
        - data_cupy (cupy.ndarray): data signal, 2-dimensional, already on GPU
        - mask_pixels_available_centrosymmetric_cupy (cupy.ndarray): mask with the available pixels to do centrosymmetric caving (already on GPU)
        - x_shifted_cupy (cupy.ndarray): shifted vector of index along the first dimension of data array, already on GPU
        - y_shifted_cupy (cupy.ndarray): shifted vector of index along the second dimension of data array, already on GPU
        - Dummy (int): pixels in data with dummy value will be replaced
        - mask_static_cupy (cupy.ndarray): mask with the pixels that will be replaced, already on GPU
        - flip_caving (bool): if True, adds horizontal-symmetric and vertical-symmetric caving (by default, centro-symmetric caving is always performed)
        - flip_horizontally_preference (bool): if flip_caving is activated, horizontal caving is done before vertical caving
        - mask_pixels_available_horizontal_cupy (cupy.ndarray): mask with the available pixels to do horizontal-symmetric caving (already on GPU)
        - mask_pixels_available_vertical_cupy (cupy.ndarray): mask with the available pixels to do vertical-symmetric caving (already on GPU)
        - x_vector_cupy (cupy.ndarray): vector of index along the first dimension of data array, already on GPU
        - y_vector_cupy (cupy.ndarray): vector of index along the second dimension of data array, already on GPU
    Outputs:
        - cupy.ndarray: data with replaced pixels (caved), still on GPU
    """
    mask_to_cave = _mask_to_cave(data_cupy, Dummy, mask_static_cupy)
    data_cupy = _process_data_caving(
        data=data_cupy,
        Dummy=Dummy,
        mask_dummy_pixels=True,
        pixels_to_cave=mask_to_cave,
        pixels_to_substitute=mask_pixels_available_centrosymmetric_cupy,
        x_vector=x_shifted_cupy,
        y_vector=y_shifted_cupy,
        use_cupy=True,
    )
    if flip_caving and Dummy:
        if flip_horizontally_preference:
            # First horizontal caving
            mask_to_cave = data_cupy == Dummy
            data_cupy = _process_data_caving(
                data=data_cupy,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_horizontal_cupy,
                x_vector=x_shifted_cupy,
                y_vector=y_vector_cupy,
                use_cupy=True,
            )
            # Second vertical caving
            mask_to_cave = data_cupy == Dummy
            data_cupy = _process_data_caving(
                data=data_cupy,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_vertical_cupy,
                x_vector=x_vector_cupy,
                y_vector=y_shifted_cupy,
                use_cupy=True,
            )
        else:
            # First vertical caving
            mask_to_cave = data_cupy == Dummy
            data_cupy = _process_data_caving(
                data=data_cupy,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_vertical_cupy,
                x_vector=x_vector_cupy,
                y_vector=y_shifted_cupy,
                use_cupy=True,
            )
            # Second horizontal caving
            mask_to_cave = data_cupy == Dummy
            data_cupy = _process_data_caving(
                data=data_cupy,
                Dummy=Dummy,
                mask_dummy_pixels=False,
                pixels_to_cave=mask_to_cave,
                pixels_to_substitute=mask_pixels_available_horizontal_cupy,
                x_vector=x_shifted_cupy,
                y_vector=y_vector_cupy,
                use_cupy=True,
            )
    return data_cupy


def process_data_caving(
    data: numpy.ndarray,
    Center_1: int,
    Center_2: int,
    Dummy: int = None,
    filename_mask_static: str = None,
    filename_mask_reference: str = None,
    algorithm_cave: str = "numpy",
    flip_caving: bool = False,
    flip_horizontally_preference: bool = True,
    **kwargs,
) -> numpy.ndarray:
    """Performs the caving on a 3-dimensional dataset using numpy methods

    Inputs:
        - data (numpy.ndarray): data signal, 2 or 3-dimensional
        - Center_1 (int): beam center in the first data dimension
        - Center_2 (int): beam center in the second data dimension
        - Dummy (int): pixels in data with dummy value will be replaced
        - filename_mask_static (str): path to the file with the mask whose pixels we want to replace
        - filename_mask_reference (str): path to the file with the mask whose pixels should not be used, that means that their symmetric brother should not be caved
        - algorithm (str): implementation to perform the caving: numpy or cupy
        - flip_caving (bool): if True, adds horizontal-symmetric and vertical-symmetric caving (by default, centro-symmetric caving is always performed)
        - flip_horizontally_preference (bool): if flip_caving is activated, horizontal caving is done before vertical caving
        - **kwargs
    Outputs:
        - numpy.ndarray: data with replaced pixels (caved) using numpy methods
    """
    if data.ndim == 2:
        data_signal_shape = data.shape
        data = data[numpy.newaxis, ...]
    elif data.ndim == 3:
        data_signal_shape = data.shape[1:]

    if algorithm_cave not in ALGORITHMS_AVAILABLE:
        logger.warning(
            f"Algorithm '{algorithm_cave}' is not available. Using '{DEFAULT_ALGORITHM}' instead."
        )
        algorithm_cave = DEFAULT_ALGORITHM
    elif algorithm_cave == "cupy" and not CUPY_AVAILABLE:
        logger.warning(f"CuPy is not available. Using {DEFAULT_ALGORITHM} instead.")
        algorithm_cave = DEFAULT_ALGORITHM
    binning = kwargs.get("binning")
    use_cupy = True if algorithm_cave == "cupy" else False

    # Load static masks
    mask_static = None
    if filename_mask_static:
        mask_static = get_persistent_array_mask(
            filename_mask=filename_mask_static,
            data_signal_shape=data_signal_shape,
            binning=binning,
            use_cupy=use_cupy,
        )
    if mask_static is not None:
        mask_static = mask_static.astype(bool)
    mask_reference = None
    if filename_mask_reference:
        mask_reference = get_persistent_array_mask(
            filename_mask=filename_mask_reference,
            data_signal_shape=data_signal_shape,
            binning=binning,
            use_cupy=False,  # This reference mask won't be cupy
        )
    if mask_reference is not None:
        mask_reference = mask_reference.astype(bool)

    # This is the mask that contains the only pixels which can be substituted, limited by the detector limits and the reference mask (it's data independent)
    mask_pixels_available_centrosymmetric = _mask_caving(
        data_signal_shape,
        Center_1,
        Center_2,
        mask_reference,
        vertical_symmetry=True,
        horizontal_symmetry=True,
        use_cupy=use_cupy,
    )
    if flip_caving:
        mask_pixels_available_horizontal = _mask_caving(
            data_signal_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=False,
            horizontal_symmetry=True,
            use_cupy=use_cupy,
        )
        mask_pixels_available_vertical = _mask_caving(
            data_signal_shape,
            Center_1,
            Center_2,
            mask_reference,
            vertical_symmetry=True,
            horizontal_symmetry=False,
            use_cupy=use_cupy,
        )
    else:
        mask_pixels_available_horizontal = None
        mask_pixels_available_vertical = None

    params_caving = {
        "dataset": data,
        "Center_1": Center_1,
        "Center_2": Center_2,
        "Dummy": Dummy,
        "mask_static": mask_static,
        "mask_reference": mask_reference,
        "flip_caving": flip_caving,
        "flip_horizontally_preference": flip_horizontally_preference,
        "mask_pixels_available_centrosymmetric": mask_pixels_available_centrosymmetric,
        "mask_pixels_available_horizontal": mask_pixels_available_horizontal,
        "mask_pixels_available_vertical": mask_pixels_available_vertical,
    }
    result = ALGORITHMS_AVAILABLE[algorithm_cave]["algorithm"](
        **params_caving,
    )
    return result


ALGORITHMS_AVAILABLE = {
    "numpy": {"algorithm": _process_dataset_caving_numpy, "use_cupy": False},
    "cupy": {"algorithm": _process_dataset_caving_cupy, "use_cupy": True},
}
DEFAULT_ALGORITHM = "numpy"
