"""This file contains functions for loading and preprocessing pianoroll data.
"""
import logging
import numpy as np
import tensorflow.compat.v1 as tf
from musegan.config import SHUFFLE_BUFFER_SIZE, PREFETCH_SIZE
LOGGER = logging.getLogger(__name__)

# --- Data loader --------------------------------------------------------------
def load_data_from_npy(filename):
    """Load and return the training data from a npy file."""
    return np.load(filename)

def load_data_from_npz(filename):
    """Load and return the training data from a npz file (sparse format)."""
    with np.load(filename) as f:
        data = np.zeros(f['shape'], np.bool_)
        data[[x for x in f['nonzero']]] = True
    return data

def load_data(data_source, data_filename):
    """Load and return the training data."""
    if data_source == 'sa':
        import SharedArray as sa
        return sa.attach(data_filename)
    if data_source == 'npy':
        return load_data_from_npy(data_filename)
    if data_source == 'npz':
        return load_data_from_npz(data_filename)
    raise ValueError("Expect `data_source` to be one of 'sa', 'npy', 'npz'. "
                     "But get " + str(data_source))

# --- Dataset Utilities -------------------------------------------------------
def random_transpose(pianoroll):
    """Randomly transpose a pianoroll with [-5, 6] semitones."""
    semitone = np.random.randint(-5, 6)
    if semitone > 0:
        pianoroll[:, semitone:, 1:] = pianoroll[:, :-semitone, 1:]
        pianoroll[:, :semitone, 1:] = 0
    elif semitone < 0:
        pianoroll[:, :semitone, 1:] = pianoroll[:, -semitone:, 1:]
        pianoroll[:, semitone:, 1:] = 0
    return pianoroll

def set_pianoroll_shape(pianoroll, data_shape):
    """Set the pianoroll shape and return the pianoroll."""
    pianoroll.set_shape(data_shape)
    return pianoroll

def set_label_shape(label):
    """Set the label shape and return the label."""
    label.set_shape([1])
    return label

# --- Sampler ------------------------------------------------------------------
def get_samples(n_samples, data, labels=None, use_random_transpose=False):
    """Return some random samples of the training data."""
    indices = np.random.choice(len(data), n_samples, False)
    if np.issubdtype(data.dtype, np.bool_):
        sample_data = data[indices] * 2. - 1.
    else:
        sample_data = data[indices]
    if use_random_transpose:
        sample_data = np.array([random_transpose(x) for x in sample_data])
    if labels is None:
        return sample_data
    return sample_data, labels[indices]

# --- Tensorflow Dataset -------------------------------------------------------
def _gen_data(data, labels=None):
    """Data Generator."""
    if labels is None:
        for item in data:
            if np.issubdtype(data.dtype, np.bool_):
                yield item * 2. - 1.
            else:
                yield item
    else:
        for i, item in enumerate(data):
            if np.issubdtype(data.dtype, np.bool_):
                yield (item * 2. - 1., labels[i])
            else:
                yield (item, labels[i])

def get_dataset(data, labels=None, batch_size=None, data_shape=None,
                use_random_transpose=False, num_threads=1):
    """Create  and return a tensorflow dataset from an array."""
    if labels is None:
        dataset = tf.data.Dataset.from_generator(
            lambda: _gen_data(data), tf.float32)
        if use_random_transpose:
            dataset = dataset.map(
                lambda pianoroll: tf.py_func(
                    random_transpose, [pianoroll], tf.float32),
                num_parallel_calls=num_threads)
        dataset = dataset.map(lambda pianoroll: set_pianoroll_shape(
            pianoroll, data_shape), num_parallel_calls=num_threads)
    else:
        assert len(data) == len(labels), (
            "Lengths of `data` and `lables` do not match.")
        dataset = tf.data.Dataset.from_generator(
            lambda: _gen_data(data, labels), [tf.float32, tf.int32])
        if use_random_transpose:
            dataset = dataset.map(
                lambda pianoroll, label: (
                    tf.py_func(random_transpose, [pianoroll], tf.float32),
                    label),
                num_parallel_calls=num_threads)
        dataset = dataset.map(
            lambda pianoroll, label: (set_pianoroll_shape(
                pianoroll, data_shape), set_label_shape(label)),
            num_parallel_calls=num_threads)

    dataset = dataset.shuffle(SHUFFLE_BUFFER_SIZE).repeat().batch(batch_size)
    return dataset.prefetch(PREFETCH_SIZE)

[((7, 9, 7, 36), 'logging.getLogger', 'logging.getLogger', ({(7, 27, 7, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((12, 11, 12, 28), 'numpy.load', 'np.load', ({(12, 19, 12, 27): 'filename'}, {}), '(filename)', True, 'import numpy as np\n'), ((36, 15, 36, 39), 'numpy.random.randint', 'np.random.randint', ({(36, 33, 36, 35): '-5', (36, 37, 36, 38): '6'}, {}), '(-5, 6)', True, 'import numpy as np\n'), ((59, 7, 59, 42), 'numpy.issubdtype', 'np.issubdtype', ({(59, 21, 59, 31): 'data.dtype', (59, 33, 59, 41): 'np.bool_'}, {}), '(data.dtype, np.bool_)', True, 'import numpy as np\n'), ((16, 9, 16, 26), 'numpy.load', 'np.load', ({(16, 17, 16, 25): 'filename'}, {}), '(filename)', True, 'import numpy as np\n'), ((17, 15, 17, 45), 'numpy.zeros', 'np.zeros', ({(17, 24, 17, 34): "f['shape']", (17, 36, 17, 44): 'np.bool_'}, {}), "(f['shape'], np.bool_)", True, 'import numpy as np\n'), ((25, 15, 25, 39), 'SharedArray.attach', 'sa.attach', ({(25, 25, 25, 38): 'data_filename'}, {}), '(data_filename)', True, 'import SharedArray as sa\n'), ((74, 15, 74, 50), 'numpy.issubdtype', 'np.issubdtype', ({(74, 29, 74, 39): 'data.dtype', (74, 41, 74, 49): 'np.bool_'}, {}), '(data.dtype, np.bool_)', True, 'import numpy as np\n'), ((80, 15, 80, 50), 'numpy.issubdtype', 'np.issubdtype', ({(80, 29, 80, 39): 'data.dtype', (80, 41, 80, 49): 'np.bool_'}, {}), '(data.dtype, np.bool_)', True, 'import numpy as np\n'), ((93, 34, 94, 62), 'tensorflow.compat.v1.py_func', 'tf.py_func', ({(94, 20, 94, 36): 'random_transpose', (94, 38, 94, 49): '[pianoroll]', (94, 51, 94, 61): 'tf.float32'}, {}), '(random_transpose, [pianoroll], tf.float32)', True, 'import tensorflow.compat.v1 as tf\n'), ((106, 20, 106, 73), 'tensorflow.compat.v1.py_func', 'tf.py_func', ({(106, 31, 106, 47): 'random_transpose', (106, 49, 106, 60): '[pianoroll]', (106, 62, 106, 72): 'tf.float32'}, {}), '(random_transpose, [pianoroll], tf.float32)', True, 'import tensorflow.compat.v1 as tf\n')]