# Path: kwiklib/dataio/selection.py
# def select(data, indices=None):
#     """Select portion of the data, with the only assumption that indices are
#     along the first axis.
#     
#     data can be a NumPy or Pandas object.
#     
#     """
#     # indices=None or 'all' means select all.
#     if indices is None or indices is 'all':
#         if type(data) == tuple:
#             indices = np.ones(data[0].shape[0], dtype=np.bool)
#         else:
#             return data
#         
#     if not hasattr(indices, '__len__') and not isinstance(indices, slice):
#         indices = [indices]
#     indices_argument = indices
#         
#     # Ensure indices is an array of indices or boolean masks.
#     if not isinstance(indices, np.ndarray) and not isinstance(indices, slice):
#         # Deal with empty indices.
#         if not len(indices):
#             if data.ndim == 1:
#                 return np.array([])
#             elif data.ndim == 2:
#                 return np.array([[]])
#             elif data.ndim == 3:
#                 return np.array([[[]]])
#         else:
#             if type(indices[0]) in (int, np.int32, np.int64):
#                 indices = np.array(indices, dtype=np.int32)
#             elif type(indices[0]) == bool:
#                 indices = np.array(indices, dtype=np.bool)
#             else:
#                 indices = np.array(indices)
#     
#     # Use NumPy, PyTables (tuple) or Pandas version
#     if type(data) == np.ndarray:
#         if data.size == 0:
#             return data
#         return select_numpy(data, indices_argument)
#     elif type(data) in (tuple, tb.EArray):
#         if type(data) == tuple:
#             data, process_fun = data
#         else:
#             process_fun = None
#         return select_pytables(data, indices_argument,
#                                process_fun=process_fun)
#     elif hasattr(data, 'values'):
#         if data.values.size == 0:
#             return data
#         return select_pandas(data, indices_argument)
#     else:
#         return select_pytables(data, indices_argument)
# 
# def slice_to_indices(indices, stop=None, lenindices=None, keys=None):
#     start, step = (indices.start or 0), (indices.step or 1)
#     if not stop:
#         assert lenindices is not None
#         # Infer stop such that indices and values have the same size.
#         stop = np.floor(start + step*lenindices)
#     indices = np.arange(start, stop, step).astype(np.int32)
#     if keys is not None:
#         indices = np.array(keys)[indices]
#     if not lenindices:
#         lenindices = len(indices)
#     assert len(indices) == lenindices
#     return indices


from collections import Counter
from kwiklib.dataio.selection import select, slice_to_indices
import tempfile
import os
import numpy as np
import pandas as pd
import tables as tb



# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def generate_clusters(indices, nspikes=100):
    """Generate all spikes in cluster 0, except some in cluster 1."""
    # 2 different clusters, with 3 spikes in cluster 1
    clusters = np.zeros(nspikes, dtype=np.int32)
    clusters[indices] = 1
    return clusters
    
def generate_earray(a):
    dirpath = tempfile.mkdtemp()
    path = os.path.join(dirpath, 'test.h5')
    f = tb.openFile(path, 'w')
    arr = f.createEArray('/', 'arr', tb.Float32Atom(), (0, a.shape[1]))
    arr.append(a)
    f.flush()
    f.close()
    return path
    
def generate_data2D(nspikes=100, ncols=5):
    data = np.random.randn(nspikes, ncols)
    return data

def assert_indexing(indices, n=10):
    values = range(n)
    stop = indices.stop or n

# Path: kwiklib/dataio/selection.py
# def select(data, indices=None):
#     """Select portion of the data, with the only assumption that indices are
#     along the first axis.
#     
#     data can be a NumPy or Pandas object.
#     
#     """
#     # indices=None or 'all' means select all.
#     if indices is None or indices is 'all':
#         if type(data) == tuple:
#             indices = np.ones(data[0].shape[0], dtype=np.bool)
#         else:
#             return data
#         
#     if not hasattr(indices, '__len__') and not isinstance(indices, slice):
#         indices = [indices]
#     indices_argument = indices
#         
#     # Ensure indices is an array of indices or boolean masks.
#     if not isinstance(indices, np.ndarray) and not isinstance(indices, slice):
#         # Deal with empty indices.
#         if not len(indices):
#             if data.ndim == 1:
#                 return np.array([])
#             elif data.ndim == 2:
#                 return np.array([[]])
#             elif data.ndim == 3:
#                 return np.array([[[]]])
#         else:
#             if type(indices[0]) in (int, np.int32, np.int64):
#                 indices = np.array(indices, dtype=np.int32)
#             elif type(indices[0]) == bool:
#                 indices = np.array(indices, dtype=np.bool)
#             else:
#                 indices = np.array(indices)
#     
#     # Use NumPy, PyTables (tuple) or Pandas version
#     if type(data) == np.ndarray:
#         if data.size == 0:
#             return data
#         return select_numpy(data, indices_argument)
#     elif type(data) in (tuple, tb.EArray):
#         if type(data) == tuple:
#             data, process_fun = data
#         else:
#             process_fun = None
#         return select_pytables(data, indices_argument,
#                                process_fun=process_fun)
#     elif hasattr(data, 'values'):
#         if data.values.size == 0:
#             return data
#         return select_pandas(data, indices_argument)
#     else:
#         return select_pytables(data, indices_argument)
# 
# def slice_to_indices(indices, stop=None, lenindices=None, keys=None):
#     start, step = (indices.start or 0), (indices.step or 1)
#     if not stop:
#         assert lenindices is not None
#         # Infer stop such that indices and values have the same size.
#         stop = np.floor(start + step*lenindices)
#     indices = np.arange(start, stop, step).astype(np.int32)
#     if keys is not None:
#         indices = np.array(keys)[indices]
#     if not lenindices:
#         lenindices = len(indices)
#     assert len(indices) == lenindices
#     return indices


# Path: kwiklib/dataio/tests/test_selection.py
from collections import Counter
from kwiklib.dataio.selection import select, slice_to_indices
import tempfile
import os
import numpy as np
import pandas as pd
import tables as tb




# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def generate_clusters(indices, nspikes=100):
    """Generate all spikes in cluster 0, except some in cluster 1."""
    # 2 different clusters, with 3 spikes in cluster 1
    clusters = np.zeros(nspikes, dtype=np.int32)
    clusters[indices] = 1
    return clusters
    
def generate_earray(a):
    dirpath = tempfile.mkdtemp()
    path = os.path.join(dirpath, 'test.h5')
    f = tb.openFile(path, 'w')
    arr = f.createEArray('/', 'arr', tb.Float32Atom(), (0, a.shape[1]))
    arr.append(a)
    f.flush()
    f.close()
    return path
    
def generate_data2D(nspikes=100, ncols=5):
    data = np.random.randn(nspikes, ncols)
    return data

def assert_indexing(indices, n=10):
    values = range(n)
    stop = indices.stop or n

    assert np.array_equal(slice_to_indices(indices, stop=stop), 