import math


def _compute_optimal_chunk_size(chunk_pixels, data_chunksize, data_shape, n_chunks_min=4):
    """
    Compute the best chunk size for the 'data' array based on existing size and
    chunk size of the `data` array and the desired number of pixels in the chunk.
    The new chunk size will contain whole number of original chunks of the `data`
    array and at least `chunk_pixels` pixels. Image pixels are determined by axes
    0 and 1. The remaining axes (typically axis 2) are not considered during this
    rechunking procedure.

    Parameters
    ----------
    chunk_pixels: int
        The desired number of pixels in the new chunks
    data_chunksize: tuple(int)
        (chunk_y, chunk_x) - original chunk size of the data array
    data_shape: tuple(int)
        (ny, nx) - the shape of the data array
    n_chunks_min: int
        The minimum number of chunks.

    Returns
    -------
    (chunk_y, chunk_x): tuple(int)
        The size of the new chunks along axes 0 and 1
    """

    if not isinstance(data_chunksize, tuple) or len(data_chunksize) != 2:
        raise ValueError(f"Unsupported value of parameter 'data_chunksize': {data_chunksize}")
    if not isinstance(data_shape, tuple) or len(data_shape) != 2:
        raise ValueError(f"Unsupported value of parameter 'data_shape': {data_shape}")

    # Compute chunk size based on the desired number of pixels and chunks in the 'data' array
    dset_chunks, dset_shape = data_chunksize, data_shape
    # Desired number of pixels in the chunk
    n_pixels_total = dset_shape[0] * dset_shape[1]
    if n_pixels_total > n_chunks_min:
        # We want to have at least 1 pixel in the chunk
        chunk_pixels = min(chunk_pixels, n_pixels_total // n_chunks_min)
    # Try to make chunks square (math.sqrt)
    chunk_x = int(math.ceil(math.sqrt(chunk_pixels) / dset_chunks[1]) * dset_chunks[1])
    chunk_x = min(chunk_x, dset_shape[1])  # The array may have few columns
    chunk_y = int(math.ceil(chunk_pixels / chunk_x / dset_chunks[0]) * dset_chunks[0])
    if chunk_y > dset_shape[0]:
        # Now explore the case if the array has small number of rows. We may want to stretch
        #   the chunk horizontally
        chunk_y = dset_shape[0]
        chunk_x = int(math.ceil(chunk_pixels / chunk_y / dset_chunks[1]) * dset_chunks[1])
        chunk_x = min(chunk_x, dset_shape[1])

    return chunk_y, chunk_x