"""Tests for `_util`."""


import numpy as np
import pytest
from numpy.testing import assert_array_equal

from skimage.morphology import _util


@pytest.mark.parametrize("image_shape", [
    (111,), (33, 44), (22, 55, 11), (6, 5, 4, 3)
])
@pytest.mark.parametrize("order", ["C", "F"])
def test_offsets_to_raveled_neighbors_highest_connectivity(image_shape, order):
    """
    Check scenarios where footprint is always of the highest connectivity
    and all dimensions are > 2.
    """
    footprint = np.ones((3,) * len(image_shape), dtype=bool)
    center = (1,) * len(image_shape)
    offsets = _util._offsets_to_raveled_neighbors(
        image_shape, footprint, center, order
    )

    # Assert only neighbors are present, center was removed
    assert len(offsets) == footprint.sum() - 1
    assert 0 not in offsets
    # Assert uniqueness
    assert len(set(offsets)) == offsets.size
    # offsets form pairs of with same value but different signs
    # if footprint is symmetric around center
    assert all(-x in offsets for x in offsets)

    # Construct image whose values are the Manhattan distance to its center
    image_center = tuple(s // 2 for s in image_shape)
    coords = [
        np.abs(np.arange(s, dtype=np.intp) - c)
        for s, c in zip(image_shape, image_center)
    ]
    grid = np.meshgrid(*coords, indexing="ij")
    image = np.sum(grid, axis=0)

    image_raveled = image.ravel(order)
    image_center_raveled = np.ravel_multi_index(
        image_center, image_shape, order=order
    )

    # Sample raveled image around its center
    samples = []
    for offset in offsets:
        index = image_center_raveled + offset
        samples.append(image_raveled[index])

    # Assert that center with value 0 wasn't selected
    assert np.min(samples) == 1
    # Assert that only neighbors where selected
    # (highest value == connectivity)
    assert np.max(samples) == len(image_shape)
    # Assert that nearest neighbors are selected first
    assert list(sorted(samples)) == samples


@pytest.mark.parametrize("image_shape", [
    (2,), (2, 2), (2, 1, 2), (2, 2, 1, 2), (0, 2, 1, 2)
])
@pytest.mark.parametrize("order", ["C", "F"])
def test_offsets_to_raveled_neighbors_footprint_smaller_image(image_shape,
                                                              order):
    """
    Test if a dimension indicated by `image_shape` is smaller than in
    `footprint`.
    """
    footprint = np.ones((3,) * len(image_shape), dtype=bool)
    center = (1,) * len(image_shape)
    offsets = _util._offsets_to_raveled_neighbors(
        image_shape, footprint, center, order
    )

    # Assert only neighbors are present, center and duplicates (possible
    # for this scenario) where removed
    assert len(offsets) <= footprint.sum() - 1
    assert 0 not in offsets
    # Assert uniqueness
    assert len(set(offsets)) == offsets.size
    # offsets form pairs of with same value but different signs
    # if footprint is symmetric around center
    assert all(-x in offsets for x in offsets)


def test_offsets_to_raveled_neighbors_explicit_0():
    """Check reviewed example."""
    image_shape = (100, 200, 3)
    footprint = np.ones((3, 3, 3), dtype=bool)
    center = (1, 1, 1)
    offsets = _util._offsets_to_raveled_neighbors(
        image_shape, footprint, center
    )

    desired = np.array([
        -600,   -3,   -1,    1,    3,  600, -603, -601, -599, -597,   -4,
         -2,    2,    4,  597,  599,  601,  603, -604, -602, -598, -596,
        596,  598,  602,  604
    ])
    assert_array_equal(offsets, desired)


def test_offsets_to_raveled_neighbors_explicit_1():
    """Check reviewed example where footprint is larger in last dimension."""
    image_shape = (10, 9, 8, 3)
    footprint = np.ones((3, 3, 3, 4), dtype=bool)
    center = (1, 1, 1, 1)
    offsets = _util._offsets_to_raveled_neighbors(
        image_shape, footprint, center
    )

    desired = np.array([
        -216,  -24,   -3,   -1,    1,    3,   24,  216, -240, -219, -217, -215, -213,
        -192,  -27,  -25,  -23,  -21,   -4,   -2,    2,    4,   21,   23,   25,   27,
         192,  213,  215,  217,  219,  240, -243, -241, -239, -237, -220, -218, -214,
        -212, -195, -193, -191, -189,  -28,  -26,  -22,  -20,   20,   22,   26,   28,
         189,  191,  193,  195,  212,  214,  218,  220,  237,  239,  241,  243, -244,
        -242, -238, -236, -196, -194, -190, -188,  188,  190,  194,  196,  236,  238,
         242,  244,    5, -211,  -19,   29,  221, -235, -187,  197,  245
    ])
    assert_array_equal(offsets, desired)
