from fullconvnets import utils
from skimage._shared import testing
from skimage._shared.testing import assert_equal, assert_array_equal

import numpy as np


def test_montage_3d_simple_gray():
    n_images, n_planes, n_rows, n_cols = 4, 2, 3, 4
    array_input = np.arange(n_images * n_planes * n_rows * n_cols,
                            dtype=np.float)
    array_input = array_input.reshape(n_images, n_planes, n_rows, n_cols)

    array_output = utils.montage_3d(array_input)
    print(array_output)
    array_reference = np.array(
        [[[ 0.,   1.,   2.,   3.,  47.5, 47.5, 47.5, 47.5],
          [ 4.,   5.,   6.,   7.,  47.5, 47.5, 47.5, 47.5],
          [ 8.,  9.,  10.,  11.,  47.5, 47.5, 47.5, 47.5],
          [48.,  49.,  50.,  51.,  47.5, 47.5, 47.5, 47.5],
          [52.,  53.,  54.,  55.,  47.5, 47.5, 47.5, 47.5],
          [56.,  57.,  58.,  59.,  47.5, 47.5, 47.5, 47.5]],

         [[12.,  13.,  14.,  15.,  47.5, 47.5, 47.5, 47.5],
          [16.,  17.,  18.,  19.,  47.5, 47.5, 47.5, 47.5],
          [20.,  21.,  22.,  23.,  47.5, 47.5, 47.5, 47.5],
          [60.,  61.,  62.,  63.,  47.5, 47.5, 47.5, 47.5],
          [64.,  65.,  66.,  67.,  47.5, 47.5, 47.5, 47.5],
          [68.,  69.,  70.,  71.,  47.5, 47.5, 47.5, 47.5]],

         [[47.5, 47.5, 47.5, 47.5, 24.,  25.,  26.,  27. ],
          [47.5, 47.5, 47.5, 47.5, 28.,  29.,  30.,  31. ],
          [47.5, 47.5, 47.5, 47.5, 32.,  33.,  34.,  35. ],
          [47.5, 47.5, 47.5, 47.5, 72.,  73.,  74.,  75. ],
          [47.5, 47.5, 47.5, 47.5, 76.,  77.,  78.,  79. ],
          [47.5, 47.5, 47.5, 47.5, 80.,  81.,  82.,  83. ]],

         [[47.5, 47.5, 47.5, 47.5, 36.,  37.,  38.,  39. ],
          [47.5, 47.5, 47.5, 47.5, 40.,  41.,  42.,  43. ],
          [47.5, 47.5, 47.5, 47.5, 44.,  45.,  46.,  47. ],
          [47.5, 47.5, 47.5, 47.5, 84.,  85.,  86.,  87. ],
          [47.5, 47.5, 47.5, 47.5, 88.,  89.,  90.,  91. ],
          [47.5, 47.5, 47.5, 47.5, 92.,  93.,  94.,  95. ]]])
    assert_array_equal(array_output, array_reference)
