from typing import Iterable


def manage_advanced_indexing(idx, single_element_getter, max_length,
                             collate_fn):
    """
    Utility function used to manage the advanced indexing and slicing.

    If more than a pattern is selected, the X and Y values will be merged
    in two separate torch Tensor objects using the stack operation.

    :param idx: Either an in, a slice object or a list (including ndarrays and
        torch Tensors) of indexes.
    :param single_element_getter: A callable used to obtain a single element
        given its int index.
    :param max_length: The maximum sequence length.
    :param collate_fn: The function to use to create a batch of data from
        single elements.
    :return: A tuple consisting of two tensors containing the X and Y values
        of the patterns addressed by the idx parameter.
    """
    indexes_iterator: Iterable[int]

    # Makes dataset sliceable
    if isinstance(idx, slice):
        indexes_iterator = range(*idx.indices(max_length))
    elif isinstance(idx, int):
        indexes_iterator = [idx]
    elif hasattr(idx, 'shape') and len(getattr(idx, 'shape')) == 0:
        # Manages 0-d ndarray / Tensor
        indexes_iterator = [int(idx)]
    else:
        indexes_iterator = idx

    elements = []
    for single_idx in indexes_iterator:
        single_element = single_element_getter(int(single_idx))
        elements.append(single_element)

    if len(elements) == 1:
        return elements[0]

    return collate_fn(elements)