import bz2
import pickle


def save_dataset_dump(input_space, labels, feature_names, label_names, filename=None):
    """Saves a compressed data set dump

    Parameters
    ----------
    input_space: array-like of array-likes
        Input space array-like of input feature vectors
    labels: array-like of binary label vectors
        Array-like of labels assigned to each input vector, as a binary
        indicator vector (i.e. if 5th position has value 1
        then the input vector has label no. 5)
    feature_names: array-like,optional
        names of features
    label_names: array-like, optional
        names of labels
    filename : str, optional
        Path to dump file, if without .bz2, the .bz2 extension will be
        appended.
    """
    data = {'X': input_space, 'y': labels, 'features': feature_names, 'labels': label_names}

    if filename is not None:
        if filename[-4:] != '.bz2':
            filename += ".bz2"

        with bz2.BZ2File(filename, "wb") as file_handle:
            pickle.dump(data, file_handle)
    else:
        return data