# Copyright (c) 2022 CNES
#
# All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.
import numpy
import pytest
import zarr

from .. import indexing


def test_slice_length():
    """Test the length of a slice."""
    assert indexing.slice_length(slice(0, 10), 10) == 10
    assert indexing.slice_length(slice(0, 10, 2), 10) == 5
    assert indexing.slice_length(slice(0, 10, 3), 10) == 4
    assert indexing.slice_length(slice(0, 10, 4), 10) == 3
    assert indexing.slice_length(slice(0, 10, 5), 10) == 2
    assert indexing.slice_length(slice(0, 10, 6), 10) == 2
    assert indexing.slice_length(slice(0, 10, 7), 10) == 2
    assert indexing.slice_length(slice(0, 10, 8), 10) == 2
    assert indexing.slice_length(slice(0, 10, 9), 10) == 2
    assert indexing.slice_length(slice(0, 10, 10), 10) == 1
    assert indexing.slice_length(slice(None), 10) == 10
    assert indexing.slice_length(slice(None, None), 10) == 10
    assert indexing.slice_length(slice(None, None, None), 10) == 10
    assert indexing.slice_length(slice(None, None, 2), 10) == 5
    assert indexing.slice_length(slice(None, None, 3), 10) == 4
    assert indexing.slice_length(slice(None, None, 4), 10) == 3
    assert indexing.slice_length(slice(None, -5), 10) == 0
    assert indexing.slice_length(slice(None, -5, 2), 10) == 0
    assert indexing.slice_length(slice(None, -5, 3), 10) == 0
    assert indexing.slice_length(slice(None, -5, 4), 10) == 0
    assert indexing.slice_length(slice(-8, -1), 10) == 7


def test_calculate_shape():
    """Test the shape calculation."""
    assert indexing.calculate_shape([slice(0, 10), slice(0, 10)],
                                    (10, 10)) == (10, 10)
    assert indexing.calculate_shape([slice(0, 5), slice(0, 5)],
                                    (10, 10)) == (5, 5)
    assert indexing.calculate_shape([slice(2, 5), slice(2, 5)],
                                    (10, 10)) == (3, 3)


def test_expected_shape():
    """Test the expected shape calculation."""
    assert indexing.expected_shape((10, 10, 10), 0, -1) == (-1, 10, 10)
    assert indexing.expected_shape((10, 10, 10), 1, -1) == (10, -1, 10)
    assert indexing.expected_shape((10, 10, 10), 2, -1) == (10, 10, -1)


def test_calculate_offset():
    """Test the offset calculation."""
    array = numpy.arange(100)
    assert indexing.calculate_offset(array, 0) == 0
    assert indexing.calculate_offset(array, 1) == 0
    assert indexing.calculate_offset(array, 2) == 1
    assert indexing.calculate_offset(array, 50) == 49

    with pytest.raises(IndexError):
        indexing.calculate_offset(array, 150)


def test_get_indexer_one_chunk():
    """Test the indexer calculation for a single chunk."""
    chunks = [zarr.array(numpy.arange(1000).reshape((10, 10, 10)))]
    indexer = indexing.get_indexer(0, chunks, slice(0, 10), (10, 10, 10), None)
    assert indexer.tensor_domain == (slice(0, 10,
                                           1), slice(0, 10,
                                                     1), slice(0, 10, 1))
    assert isinstance(indexer, indexing.Selection)
    assert isinstance(indexer, indexing.SingleSelection)
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][...])

    indexer = indexing.get_indexer(1, chunks, slice(0, 5), (10, 10, 10), None)
    assert indexer.tensor_domain == (slice(0, 5, 1), slice(0, 10,
                                                           1), slice(0, 10, 1))
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][:5, :, :])

    indexer = indexing.get_indexer(1, chunks, (slice(0, 5), slice(0, 5)),
                                   (10, 10, 10), None)
    assert indexer.tensor_domain == (slice(0, 5, 1), slice(0, 5,
                                                           1), slice(0, 10, 1))
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][:5, :5, :])

    indexer = indexing.get_indexer(2, chunks,
                                   (slice(0, 5), slice(0, 5), slice(0, 5)),
                                   (10, 10, 10), None)
    assert indexer.tensor_domain == (slice(0, 5, 1), slice(0, 5,
                                                           1), slice(0, 5, 1))
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][:5, :5, :5])

    indexer = indexing.get_indexer(2, chunks, (slice(50, 5)), (10, 10, 10),
                                   None)
    assert isinstance(indexer, indexing.Selection)
    assert isinstance(indexer, indexing.SingleSelection)

    indexer = indexing.get_indexer(0, chunks,
                                   (slice(0, 10), slice(0, 10), slice(0, 10)),
                                   (10, 10, 10),
                                   (slice(0, 5), slice(1, 5), slice(2, 5)))
    assert indexer.tensor_domain == (slice(0, 5, 1), slice(1, 5,
                                                           1), slice(2, 5, 1))
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][:5, 1:5, 2:5])

    with pytest.raises(IndexError):
        indexing.get_indexer(2, [zarr.array(numpy.arange(1000))],
                             (slice(0, 5), slice(0, 5), slice(0, 5)), (1000, ),
                             None)

    indexer = indexing.get_indexer(1, chunks, slice(0, 5, 2), (10, 10, 10),
                                   None)
    assert indexer.tensor_domain == (slice(0, 5, 2), slice(0, 10,
                                                           1), slice(0, 10, 1))
    assert numpy.all(indexer.getitem(chunks, 0) == chunks[0][0:5:2, :, :])

    with pytest.raises(ValueError):
        indexing.get_indexer(0, chunks, slice(0, 5, 2), (10, 10, 10), None)


def test_get_indexer_empty():
    """Test the indexer for an empty selection."""
    arr1 = zarr.array(numpy.full((20, 10), 1, dtype="int64"))
    arr2 = zarr.array(numpy.full((10, 10), 2, dtype="int64"))
    arr3 = zarr.array(numpy.full((5, 10), 3, dtype="int64"))
    arr4 = zarr.array(numpy.full((2, 10), 4, dtype="int64"))

    chunks = [arr1, arr2, arr3, arr4]

    indexer = indexing.get_indexer(0, chunks, slice(0, 0), (37, 10), None)
    assert isinstance(indexer, indexing.Selection)
    assert isinstance(indexer, indexing.EmptySelection)
    assert indexer.getitem(chunks, 0).shape == (0, 10)

    indexer = indexing.get_indexer(0, chunks, slice(10, 2), (37, 10), None)
    assert isinstance(indexer, indexing.Selection)
    assert isinstance(indexer, indexing.EmptySelection)
    assert indexer.getitem(chunks, 0).shape == (0, 10)
    assert indexer.getitem(chunks, 1).shape == (arr1.shape[0], 0)


@pytest.mark.parametrize("axis", [0, 1])
def test_get_indexer_multiple(axis):
    """Test the indexer for multiple chunks."""
    arr1 = numpy.full((20, 10), 1, dtype="int64")
    arr2 = numpy.full((10, 10), 2, dtype="int64")
    arr3 = numpy.full((5, 10), 3, dtype="int64")
    arr4 = numpy.full((2, 10), 4, dtype="int64")

    chunks = [arr1, arr2, arr3, arr4]

    if axis == 1:
        for ix, item in enumerate(chunks):
            chunks[ix] = item.transpose()

    arr = numpy.concatenate(chunks, axis=axis)
    chunks = [zarr.array(chunk) for chunk in chunks]

    for item in [
            slice(-20, -10),
            slice(15, 36),
            slice(20, 36),
            slice(35, 36),
            slice(5, 15),
            slice(None, -5),
            slice(None, 2),
            slice(None, 20),
            slice(None),
    ]:
        test_case = [item]
        test_case.insert([1, 0][axis], slice(None))
        indexer = indexing.get_indexer(axis, chunks, tuple(test_case),
                                       arr.shape, None)
        assert isinstance(indexer, indexing.Selection)
        assert isinstance(indexer, indexing.MultipleSelection)

        values = indexer.getitem(chunks, axis)
        expected_shape = indexing.expected_shape(
            arr.shape, axis, indexer.tensor_domain[axis].stop -
            indexer.tensor_domain[axis].start)
        assert values.shape == expected_shape
        assert numpy.all(values == arr[tuple(test_case)])
