import os
import sys
import unittest
import tempfile
import itertools
import numpy as np

from nose.tools import assert_equal, assert_raises, assert_true
from nibabel.testing import assert_arrays_equal
from numpy.testing import assert_array_equal

from ..array_sequence import ArraySequence, is_array_sequence, concatenate


SEQ_DATA = {}


def setup():
    global SEQ_DATA
    rng = np.random.RandomState(42)
    SEQ_DATA['rng'] = rng
    SEQ_DATA['data'] = generate_data(nb_arrays=5, common_shape=(3,), rng=rng)
    SEQ_DATA['seq'] = ArraySequence(SEQ_DATA['data'])


def generate_data(nb_arrays, common_shape, rng):
    data = [rng.rand(*(rng.randint(3, 20),) + common_shape) * 100
            for _ in range(nb_arrays)]
    return data


def check_empty_arr_seq(seq):
    assert_equal(len(seq), 0)
    assert_equal(len(seq._offsets), 0)
    assert_equal(len(seq._lengths), 0)
    # assert_equal(seq._data.ndim, 0)
    assert_equal(seq._data.ndim, 1)
    assert_true(seq.common_shape == ())


def check_arr_seq(seq, arrays):
    lengths = list(map(len, arrays))
    assert_true(is_array_sequence(seq))
    assert_equal(len(seq), len(arrays))
    assert_equal(len(seq._offsets), len(arrays))
    assert_equal(len(seq._lengths), len(arrays))
    assert_equal(seq._data.shape[1:], arrays[0].shape[1:])
    assert_equal(seq.common_shape, arrays[0].shape[1:])
    assert_arrays_equal(seq, arrays)

    # If seq is a view, then order of internal data is not guaranteed.
    if seq._is_view:
        # The only thing we can check is the _lengths.
        assert_array_equal(sorted(seq._lengths), sorted(lengths))
    else:
        seq.shrink_data()
        assert_equal(seq._data.shape[0], sum(lengths))
        assert_array_equal(seq._data, np.concatenate(arrays, axis=0))
        assert_array_equal(seq._offsets, np.r_[0, np.cumsum(lengths)[:-1]])
        assert_array_equal(seq._lengths, lengths)


def check_arr_seq_view(seq_view, seq):
    assert_true(seq_view._is_view)
    assert_true(seq_view is not seq)
    assert_true(np.may_share_memory(seq_view._data, seq._data))
    assert_true(seq_view._offsets is not seq._offsets)
    assert_true(seq_view._lengths is not seq._lengths)


class TestArraySequence(unittest.TestCase):

    def test_creating_empty_arraysequence(self):
        check_empty_arr_seq(ArraySequence())

    def test_creating_arraysequence_from_list(self):
        # Empty list
        check_empty_arr_seq(ArraySequence([]))

        # List of ndarrays.
        N = 5
        for ndim in range(1, N+1):
            common_shape = tuple([SEQ_DATA['rng'].randint(1, 10)
                                 for _ in range(ndim-1)])
            data = generate_data(nb_arrays=5, common_shape=common_shape,
                                 rng=SEQ_DATA['rng'])
            check_arr_seq(ArraySequence(data), data)

        # Force ArraySequence constructor to use buffering.
        buffer_size = 1. / 1024**2  # 1 bytes
        check_arr_seq(ArraySequence(iter(SEQ_DATA['data']), buffer_size),
                      SEQ_DATA['data'])

    def test_creating_arraysequence_from_generator(self):
        gen_1, gen_2 = itertools.tee((e for e in SEQ_DATA['data']))
        seq = ArraySequence(gen_1)
        seq_with_buffer = ArraySequence(gen_2, buffer_size=256)

        # Check buffer size effect
        assert_equal(seq_with_buffer.data.shape, seq.data.shape)
        assert_true(seq_with_buffer._buffer_size > seq._buffer_size)

        # Check generator result
        check_arr_seq(seq, SEQ_DATA['data'])
        check_arr_seq(seq_with_buffer, SEQ_DATA['data'])

        # Already consumed generator
        check_empty_arr_seq(ArraySequence(gen_1))

    def test_creating_arraysequence_from_arraysequence(self):
        seq = ArraySequence(SEQ_DATA['data'])
        check_arr_seq(ArraySequence(seq), SEQ_DATA['data'])

        # From an empty ArraySequence
        seq = ArraySequence()
        check_empty_arr_seq(ArraySequence(seq))

    def test_arraysequence_iter(self):
        assert_arrays_equal(SEQ_DATA['seq'], SEQ_DATA['data'])

        # Try iterating through a corrupted ArraySequence object.
        seq = SEQ_DATA['seq'].copy()
        seq._lengths = seq._lengths[::2]
        assert_raises(ValueError, list, seq)

    def test_arraysequence_copy(self):
        orig = SEQ_DATA['seq']
        seq = orig.copy()
        n_rows = seq.total_nb_rows
        assert_equal(n_rows, orig.total_nb_rows)
        assert_array_equal(seq._data, orig._data[:n_rows])
        assert_true(seq._data is not orig._data)
        assert_array_equal(seq._offsets, orig._offsets)
        assert_true(seq._offsets is not orig._offsets)
        assert_array_equal(seq._lengths, orig._lengths)
        assert_true(seq._lengths is not orig._lengths)
        assert_equal(seq.common_shape, orig.common_shape)

        # Taking a copy of an `ArraySequence` generated by slicing.
        # Only keep needed data.
        seq = orig[::2].copy()
        check_arr_seq(seq, SEQ_DATA['data'][::2])
        assert_true(seq._data is not orig._data)

    def test_arraysequence_append(self):
        element = generate_data(nb_arrays=1,
                                common_shape=SEQ_DATA['seq'].common_shape,
                                rng=SEQ_DATA['rng'])[0]

        # Append a new element.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.append(element)
        check_arr_seq(seq, SEQ_DATA['data'] + [element])

        # Append a list of list.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.append(element.tolist())
        check_arr_seq(seq, SEQ_DATA['data'] + [element])

        # Append to an empty ArraySequence.
        seq = ArraySequence()
        seq.append(element)
        check_arr_seq(seq, [element])

        # Append an empty array.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.append([])
        check_arr_seq(seq, SEQ_DATA['seq'])

        # Append an element with different shape.
        element = generate_data(nb_arrays=1,
                                common_shape=SEQ_DATA['seq'].common_shape*2,
                                rng=SEQ_DATA['rng'])[0]
        assert_raises(ValueError, seq.append, element)

    def test_arraysequence_extend(self):
        new_data = generate_data(nb_arrays=10,
                                 common_shape=SEQ_DATA['seq'].common_shape,
                                 rng=SEQ_DATA['rng'])

        # Extend with an empty list.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.extend([])
        check_arr_seq(seq, SEQ_DATA['data'])

        # Extend with a list of ndarrays.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.extend(new_data)
        check_arr_seq(seq, SEQ_DATA['data'] + new_data)

        # Extend with a generator.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.extend((d for d in new_data))
        check_arr_seq(seq, SEQ_DATA['data'] + new_data)

        # Extend with another `ArraySequence` object.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.extend(ArraySequence(new_data))
        check_arr_seq(seq, SEQ_DATA['data'] + new_data)

        # Extend with an `ArraySequence` view (e.g. been sliced).
        # Need to make sure we extend only the data we need.
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        seq.extend(ArraySequence(new_data)[::2])
        check_arr_seq(seq, SEQ_DATA['data'] + new_data[::2])

        # Test extending an empty ArraySequence
        seq = ArraySequence()
        seq.extend(ArraySequence())
        check_empty_arr_seq(seq)

        seq.extend(SEQ_DATA['seq'])
        check_arr_seq(seq, SEQ_DATA['data'])

        # Extend with elements of different shape.
        data = generate_data(nb_arrays=10,
                             common_shape=SEQ_DATA['seq'].common_shape*2,
                             rng=SEQ_DATA['rng'])
        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
        assert_raises(ValueError, seq.extend, data)

        # Extend after extracting some slice
        working_slice = seq[:2]
        seq.extend(ArraySequence(new_data))

    def test_arraysequence_getitem(self):
        # Get one item
        for i, e in enumerate(SEQ_DATA['seq']):
            assert_array_equal(SEQ_DATA['seq'][i], e)

        # Get all items using indexing (creates a view).
        indices = list(range(len(SEQ_DATA['seq'])))
        seq_view = SEQ_DATA['seq'][indices]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        # We took all elements so the view should match the original.
        check_arr_seq(seq_view, SEQ_DATA['seq'])

        # Get multiple items using ndarray of dtype integer.
        for dtype in [np.int8, np.int16, np.int32, np.int64]:
            seq_view = SEQ_DATA['seq'][np.array(indices, dtype=dtype)]
            check_arr_seq_view(seq_view, SEQ_DATA['seq'])
            # We took all elements so the view should match the original.
            check_arr_seq(seq_view, SEQ_DATA['seq'])

        # Get multiple items out of order (creates a view).
        SEQ_DATA['rng'].shuffle(indices)
        seq_view = SEQ_DATA['seq'][indices]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        check_arr_seq(seq_view, [SEQ_DATA['data'][i] for i in indices])

        # Get slice (this will create a view).
        seq_view = SEQ_DATA['seq'][::2]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        check_arr_seq(seq_view, SEQ_DATA['data'][::2])

        # Use advanced indexing with ndarray of data type bool.
        selection = np.array([False, True, True, False, True])
        seq_view = SEQ_DATA['seq'][selection]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        check_arr_seq(seq_view,
                      [SEQ_DATA['data'][i]
                       for i, keep in enumerate(selection) if keep])

        # Test invalid indexing
        assert_raises(TypeError, SEQ_DATA['seq'].__getitem__, 'abc')

        # Get specific columns.
        seq_view = SEQ_DATA['seq'][:, 2]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data']])

        # Combining multiple slicing and indexing operations.
        seq_view = SEQ_DATA['seq'][::-2][:, 2]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data'][::-2]])

    def test_arraysequence_setitem(self):
        # Set one item
        seq = SEQ_DATA['seq'] * 0
        for i, e in enumerate(SEQ_DATA['seq']):
            seq[i] = e

        check_arr_seq(seq, SEQ_DATA['seq'])

        # Setitem with a scalar.
        seq = SEQ_DATA['seq'].copy()
        seq[:] = 0
        assert_true(seq._data.sum() == 0)

        # Setitem with a list of ndarray.
        seq = SEQ_DATA['seq'] * 0
        seq[:] = SEQ_DATA['data']
        check_arr_seq(seq, SEQ_DATA['data'])

        # Setitem using tuple indexing.
        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
        seq[:, 0] = 0
        assert_true(seq._data[:, 0].sum() == 0)

        # Setitem using tuple indexing.
        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
        seq[range(len(seq))] = 0
        assert_true(seq._data.sum() == 0)

        # Setitem of a slice using another slice.
        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
        seq[0:4] = seq[5:9]
        check_arr_seq(seq[0:4], seq[5:9])

        # Setitem between array sequences with different number of sequences.
        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
        assert_raises(ValueError, seq.__setitem__, slice(0, 4), seq[5:10])

        # Setitem between array sequences with different amount of points.
        seq1 = ArraySequence(np.arange(10).reshape(5, 2))
        seq2 = ArraySequence(np.arange(15).reshape(5, 3))
        assert_raises(ValueError, seq1.__setitem__, slice(0, 5), seq2)

        # Setitem between array sequences with different common shape.
        seq1 = ArraySequence(np.arange(12).reshape(2, 2, 3))
        seq2 = ArraySequence(np.arange(8).reshape(2, 2, 2))
        assert_raises(ValueError, seq1.__setitem__, slice(0, 2), seq2)

        # Invalid index.
        assert_raises(TypeError, seq.__setitem__, object(), None)

    def test_arraysequence_operators(self):
        # Disable division per zero warnings.
        flags = np.seterr(divide='ignore', invalid='ignore')
        SCALARS = [42, 0.5, True, -3, 0]
        CMP_OPS = ["__eq__", "__ne__", "__lt__", "__le__", "__gt__", "__ge__"]

        seq = SEQ_DATA['seq'].copy()
        seq_int = SEQ_DATA['seq'].copy()
        seq_int._data = seq_int._data.astype(int)
        seq_bool = SEQ_DATA['seq'].copy() > 30

        ARRSEQS = [seq, seq_int, seq_bool]
        VIEWS = [seq[::2], seq_int[::2], seq_bool[::2]]

        def _test_unary(op, arrseq):
            orig = arrseq.copy()
            seq = getattr(orig, op)()
            assert_true(seq is not orig)
            check_arr_seq(seq, [getattr(d, op)() for d in orig])

        def _test_binary(op, arrseq, scalars, seqs, inplace=False):
            for scalar in scalars:
                orig = arrseq.copy()
                seq = getattr(orig, op)(scalar)
                assert_true((seq is orig) if inplace else (seq is not orig))
                check_arr_seq(seq, [getattr(e, op)(scalar) for e in arrseq])

            # Test math operators with another ArraySequence.
            for other in seqs:
                orig = arrseq.copy()
                seq = getattr(orig, op)(other)
                assert_true(seq is not SEQ_DATA['seq'])
                check_arr_seq(seq, [getattr(e1, op)(e2) for e1, e2 in zip(arrseq, other)])

            # Operations between array sequences of different lengths.
            orig = arrseq.copy()
            assert_raises(ValueError, getattr(orig, op), orig[::2])

            # Operations between array sequences with different amount of data.
            seq1 = ArraySequence(np.arange(10).reshape(5, 2))
            seq2 = ArraySequence(np.arange(15).reshape(5, 3))
            assert_raises(ValueError, getattr(seq1, op), seq2)

            # Operations between array sequences with different common shape.
            seq1 = ArraySequence(np.arange(12).reshape(2, 2, 3))
            seq2 = ArraySequence(np.arange(8).reshape(2, 2, 2))
            assert_raises(ValueError, getattr(seq1, op), seq2)


        for op in ["__add__", "__sub__", "__mul__", "__mod__",
                   "__floordiv__", "__truediv__"] + CMP_OPS:
            _test_binary(op, seq, SCALARS, ARRSEQS)
            _test_binary(op, seq_int, SCALARS, ARRSEQS)

            # Test math operators with ArraySequence views.
            _test_binary(op, seq[::2], SCALARS, VIEWS)
            _test_binary(op, seq_int[::2], SCALARS, VIEWS)

            if op in CMP_OPS:
                continue

            op = "__i{}__".format(op.strip("_"))
            _test_binary(op, seq, SCALARS, ARRSEQS, inplace=True)

            if op == "__itruediv__":
                continue  # Going to deal with it separately.

            _test_binary(op, seq_int, [42, -3, True, 0], [seq_int, seq_bool, -seq_int], inplace=True)  # int <-- int
            assert_raises(TypeError, _test_binary, op, seq_int, [0.5], [], inplace=True)  # int <-- float
            assert_raises(TypeError, _test_binary, op, seq_int, [], [seq], inplace=True)  # int <-- float

        # __pow__ : Integers to negative integer powers are not allowed.
        _test_binary("__pow__", seq, [42, -3, True, 0], [seq_int, seq_bool, -seq_int])
        _test_binary("__ipow__", seq, [42, -3, True, 0], [seq_int, seq_bool, -seq_int], inplace=True)
        assert_raises(ValueError, _test_binary, "__pow__", seq_int, [-3], [])
        assert_raises(ValueError, _test_binary, "__ipow__", seq_int, [-3], [], inplace=True)

        # __itruediv__ is only valid with float arrseq.
        for scalar in SCALARS + ARRSEQS:
            assert_raises(TypeError, getattr(seq_int.copy(), "__itruediv__"), scalar)

        # Bitwise operators
        for op in ("__lshift__", "__rshift__", "__or__", "__and__", "__xor__"):
            _test_binary(op, seq_bool, [42, -3, True, 0], [seq_int, seq_bool, -seq_int])
            assert_raises(TypeError, _test_binary, op, seq_bool, [0.5], [])
            assert_raises(TypeError, _test_binary, op, seq, [], [seq])

        # Unary operators
        for op in ["__neg__", "__abs__"]:
            _test_unary(op, seq)
            _test_unary(op, -seq)
            _test_unary(op, seq_int)
            _test_unary(op, -seq_int)

        _test_unary("__abs__", seq_bool)
        _test_unary("__invert__", seq_bool)
        assert_raises(TypeError, _test_unary, "__invert__", seq)

        # Restore flags.
        np.seterr(**flags)

    def test_arraysequence_repr(self):
        # Test that calling repr on a ArraySequence object is not falling.
        repr(SEQ_DATA['seq'])

        # Test calling repr when the number of arrays is bigger dans Numpy's
        # print option threshold.
        nb_arrays = 50
        seq = ArraySequence(generate_data(nb_arrays, common_shape=(1,),
                                          rng=SEQ_DATA['rng']))

        bkp_threshold = np.get_printoptions()['threshold']
        np.set_printoptions(threshold=nb_arrays*2)
        txt1 = repr(seq)
        np.set_printoptions(threshold=nb_arrays//2)
        txt2 = repr(seq)
        assert_true(len(txt2) < len(txt1))
        np.set_printoptions(threshold=bkp_threshold)

    def test_save_and_load_arraysequence(self):
        # Test saving and loading an empty ArraySequence.
        with tempfile.TemporaryFile(mode="w+b", suffix=".npz") as f:
            seq = ArraySequence()
            seq.save(f)
            f.seek(0, os.SEEK_SET)
            loaded_seq = ArraySequence.load(f)
            assert_array_equal(loaded_seq._data, seq._data)
            assert_array_equal(loaded_seq._offsets, seq._offsets)
            assert_array_equal(loaded_seq._lengths, seq._lengths)

        # Test saving and loading a ArraySequence.
        with tempfile.TemporaryFile(mode="w+b", suffix=".npz") as f:
            seq = SEQ_DATA['seq']
            seq.save(f)
            f.seek(0, os.SEEK_SET)
            loaded_seq = ArraySequence.load(f)
            assert_array_equal(loaded_seq._data, seq._data)
            assert_array_equal(loaded_seq._offsets, seq._offsets)
            assert_array_equal(loaded_seq._lengths, seq._lengths)

            # Make sure we can add new elements to it.
            loaded_seq.append(SEQ_DATA['data'][0])

    def test_get_data(self):
        seq_view = SEQ_DATA['seq'][::2]
        check_arr_seq_view(seq_view, SEQ_DATA['seq'])

        # We make sure the array sequence data does not
        # contain more elements than it is supposed to.
        data = seq_view.get_data()
        assert len(data) < len(seq_view._data)


def test_concatenate():
    seq = SEQ_DATA['seq'].copy()  # In case there is in-place modification.
    seqs = [seq[:, [i]] for i in range(seq.common_shape[0])]
    new_seq = concatenate(seqs, axis=1)
    seq._data += 100  # Modifying the 'seq' shouldn't change 'new_seq'.
    check_arr_seq(new_seq, SEQ_DATA['data'])
    assert_true(not new_seq._is_view)

    seq = SEQ_DATA['seq']
    seqs = [seq[:, [i]] for i in range(seq.common_shape[0])]
    new_seq = concatenate(seqs, axis=0)
    assert_true(len(new_seq), seq.common_shape[0] * len(seq))
    assert_array_equal(new_seq._data, seq._data.T.reshape((-1, 1)))
