from copy import copy

import numpy as np

from janggu.data import Array
from janggu.data import NanToNumConverter
from janggu.data import RandomOrientation
from janggu.data import RandomSignalScale
from janggu.data import ReduceDim


def test_nparr(tmpdir):
    X = Array("X", np.random.random((1000, 100)))
    y = Array('y', np.random.randint(2, size=(1000,)))

    np.testing.assert_equal(len(X), len(y))
    np.testing.assert_equal(len(X), 1000)
    np.testing.assert_equal(X.shape, (1000, 100,))
    np.testing.assert_equal(y.shape, (1000, 1))
    assert y.ndim == 2
    assert y.shape == (1000, 1)
    new_X = copy(X)


def test_reducedim():
    x_orig = np.zeros((3,1,1,2))

    np.testing.assert_equal(x_orig.ndim, 4)
    x_reduce = ReduceDim(Array('test', x_orig, conditions=["A", "B"]))
    x_reduce = ReduceDim(Array('test', x_orig, conditions=["A", "B"]), aggregator='mean')
    x_reduce = ReduceDim(Array('test', x_orig, conditions=["A", "B"]), aggregator='max')
    x_reduce = ReduceDim(Array('test', x_orig, conditions=["A", "B"]), aggregator=np.mean)
    np.testing.assert_equal(len(x_reduce), 3)
    np.testing.assert_equal(len(x_reduce), 3)
    np.testing.assert_equal(len(x_reduce), 3)
    np.testing.assert_equal(x_reduce.shape, (3,2))
    np.testing.assert_equal(x_reduce.ndim, 2)
    assert x_reduce[0].shape == (1, 2)
    assert x_reduce[:3].shape == (3, 2)
    assert x_reduce[[0,1]].shape == (2, 2)
    assert x_reduce.ndim == 2
    new_x = copy(x_reduce)
    assert x_reduce[0].shape == new_x[0].shape
    assert x_reduce.conditions == ["A", "B"]

def test_nantonumconverter():
    x_orig = np.zeros((3,1,1,2))
    x_orig[0,0,0,0] = np.nan
    arr = Array('test', x_orig, conditions=["A", "B"])
    assert np.isnan(arr[0].mean())

    x_tr = NanToNumConverter(Array('test', x_orig, conditions=["A", "B"]))
    assert x_tr[0].shape == (1, 1, 1, 2)
    assert x_tr[:3].shape == (3, 1, 1, 2)
    assert x_tr[[0,1]].shape == (2, 1, 1, 2)
    assert len(x_tr) == 3
    assert x_tr.shape == (3, 1, 1, 2)
    assert x_tr.ndim == 4
    assert not np.isnan(x_tr[0].mean())
    np.testing.assert_equal(x_tr[0], [[[[0,0]]]])
    new_x = copy(x_tr)
    assert x_tr[0].shape == new_x[0].shape
    assert x_tr.conditions == ["A", "B"]

def test_randomorientation():
    x_orig = np.zeros((3,1,1,2))

    x_tr = RandomOrientation(Array('test', x_orig, conditions=["A", "B"]))
    assert x_tr[0].shape == (1, 1, 1, 2)
    assert x_tr[:3].shape == (3, 1, 1, 2)
    assert x_tr[[0,1]].shape == (2, 1, 1, 2)
    np.testing.assert_equal(len(x_tr), 3)
    assert len(x_tr) == 3
    assert x_tr.shape == (3, 1, 1, 2)
    assert x_tr.ndim == 4
    np.testing.assert_equal(x_tr[0], [[[[0,0]]]])
    new_x = copy(x_tr)
    assert x_tr[0].shape == new_x[0].shape
    assert x_tr.conditions == ["A", "B"]


def test_randomsignalscale():
    x_orig = np.ones((3,1,1,2))

    x_tr = RandomSignalScale(Array('test', x_orig), .1)
    assert x_tr[0].shape == (1, 1, 1, 2)
    assert x_tr[:3].shape == (3, 1, 1, 2)
    assert x_tr[[0,1]].shape == (2, 1, 1, 2)
    np.testing.assert_equal(len(x_tr), 3)
    assert len(x_tr) == 3
    assert x_tr.shape == (3, 1, 1, 2)
    assert x_tr.ndim == 4
    new_x = copy(x_tr)
    assert x_tr[0].shape == new_x[0].shape
    assert x_tr.conditions == None
