from copy import copy

import numpy as np
import pytest

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


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)
    with pytest.raises(ValueError):
        ReduceDim(Array('test', x_orig, conditions=["A", "B"]), aggregator='nonsense')

    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_squeezedim():
    x_orig = np.zeros((3,1,1,2))

    np.testing.assert_equal(x_orig.ndim, 4)
    x_sq = SqueezeDim(Array('test', x_orig, conditions=["A", "B"]))

    np.testing.assert_equal(len(x_sq), 3)

    np.testing.assert_equal(x_sq.shape, (3,2))

    np.testing.assert_equal(x_sq.ndim, 2)
    assert x_sq[0].shape == (2,)
    assert x_sq[:3].shape == (3, 2)
    assert x_sq[[0,1]].shape == (2, 2)
    assert x_sq.ndim == 2
    new_x = copy(x_sq)
    assert x_sq[0].shape == new_x[0].shape
    assert x_sq.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


def test_randomshift():
    x_orig = np.zeros((1,4,1,4))
    x_orig[0, 0, 0,0] = 1
    x_orig[0, 1, 0,1] = 1
    x_orig[0, 2, 0,2] = 1
    x_orig[0, 3, 0,3] = 1

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