import numpy as np
import tensorflow as tf
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.platform import test
from tensorflow.contrib.layers.python.layers import layers as _layers

def _sparsify(array, threshold=0.5):
    array[array < threshold] = 0
    non_zero = np.where(array)
    indices = np.vstack(non_zero).T
    values = array[non_zero]
    shape = array.shape
    return indices, values, shape


class SparsifyTest(test.TestCase):
      def testSparsePartialFlatten(self):
        """Test `_inner_flatten` on `SparseTensor`s."""
        shape = [4, 3, 11, 6, 1, 3]
        np.random.seed(10301)
        random_ = np.random.rand(*shape)
        indices, values, _ = _sparsify(random_)

        for new_rank in [1, 2, 3, 4, 5]:
            expected_shape = (shape[:new_rank - 1] + [np.prod(shape[new_rank - 1:])])
            reshaped_random_ = np.reshape(random_, expected_shape)
            expected_indices, expected_values, _ = _sparsify(reshaped_random_)

            inputs_t = sparse_tensor.SparseTensor(indices, values, shape)

            flattened_t = _layers._inner_flatten(inputs_t, new_rank)

        with self.test_session() as sess:
            flattened = sess.run(flattened_t)

        np.testing.assert_array_equal(expected_indices, flattened.indices)
        np.testing.assert_array_equal(expected_values, flattened.values)
        np.testing.assert_array_equal(expected_shape, flattened.dense_shape)


t = SparsifyTest()
t.testSparsePartialFlatten()