from mock import patch
import numpy as np


def test_dataset_simple():
    from ..dataset import Dataset

    data = object()
    target = object()
    dataset = Dataset(data, target)
    assert dataset.data is data
    assert dataset.target is target


@patch('nolearn.dataset.np.load')
def test_dataset_with_filenames(load):
    from ..dataset import Dataset

    data = 'datafile'
    target = 'targetfile'
    dataset = Dataset(data, target)
    assert load.call_count == 2
    assert dataset.target is load.return_value


def test_dataset_train_test_split():
    from ..dataset import Dataset

    data = np.arange(100)
    target = np.array([0] * 50 + [1] * 50)
    dataset = Dataset(data, target)

    assert dataset.split_indices.classes.tolist() == [0, 1]
    assert dataset.split_indices.n_train == 75
    assert dataset.split_indices.n_test == 25

    X_train, X_test, y_train, y_test = dataset.train_test_split()
    assert len(X_train) == len(y_train)
    assert len(X_test) == len(y_test)


def test_dataset_scale():
    from ..dataset import Dataset

    data = np.arange(100).astype('float')
    target = np.array([0] * 100)
    dataset = Dataset(data, target)

    dataset.scale()
    assert dataset.data[0] == -1.7148160424389376
    assert dataset.data[-1] == 1.7148160424389376

[((15, 1, 15, 33), 'mock.patch', 'patch', ({(15, 7, 15, 32): '"""nolearn.dataset.np.load"""'}, {}), "('nolearn.dataset.np.load')", False, 'from mock import patch\n'), ((29, 11, 29, 25), 'numpy.arange', 'np.arange', ({(29, 21, 29, 24): '100'}, {}), '(100)', True, 'import numpy as np\n'), ((30, 13, 30, 42), 'numpy.array', 'np.array', ({(30, 22, 30, 41): '[0] * 50 + [1] * 50'}, {}), '([0] * 50 + [1] * 50)', True, 'import numpy as np\n'), ((46, 13, 46, 32), 'numpy.array', 'np.array', ({(46, 22, 46, 31): '[0] * 100'}, {}), '([0] * 100)', True, 'import numpy as np\n'), ((45, 11, 45, 25), 'numpy.arange', 'np.arange', ({(45, 21, 45, 24): '100'}, {}), '(100)', True, 'import numpy as np\n')]