def test_2d_bf():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels_bf = random_walker(data, labels, beta=90, mode="bf")
    assert (labels_bf[25:45, 40:60] == 2).all()
    full_prob_bf = random_walker(data, labels, beta=90, mode="bf", return_full_prob=True)
    assert (full_prob_bf[1, 25:45, 40:60] >= full_prob_bf[0, 25:45, 40:60]).all()
    return data, labels_bf, full_prob_bf

def test_2d_cg_mg():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels_cg_mg = random_walker(data, labels, beta=90, mode="cg_mg")
    assert (labels_cg_mg[25:45, 40:60] == 2).all()
    full_prob = random_walker(data, labels, beta=90, mode="cg_mg", return_full_prob=True)
    assert (full_prob[1, 25:45, 40:60] >= full_prob[0, 25:45, 40:60]).all()
    return data, labels_cg_mg

def test_length2_spacing():
    # If this passes without raising an exception (warnings OK), the new
    #   spacing code is working properly.
    np.random.seed(42)
    img = np.ones((10, 10)) + 0.2 * np.random.normal(size=(10, 10))
    labels = np.zeros((10, 10), dtype=np.uint8)
    labels[2, 4] = 1
    labels[6, 8] = 4
    random_walker(img, labels, spacing=(1., 2.))

def test_multispectral_2d():
    lx, ly = 70, 100
    data, labels = make_2d_syntheticdata(lx, ly)
    data2 = data.copy()
    data.shape += (1,)
    data = data.repeat(2, axis=2)   # Result should be identical
    multi_labels = random_walker(data, labels, mode='cg', multichannel=True)
    single_labels = random_walker(data2, labels, mode='cg')
    assert (multi_labels.reshape(labels.shape)[25:45, 40:60] == 2).all()
    return data, multi_labels, single_labels, labels

def test_multispectral_2d():
    lx, ly = 70, 100
    data, labels = make_2d_syntheticdata(lx, ly)
    data = data[..., np.newaxis].repeat(2, axis=-1)  # Expect identical output
    multi_labels = random_walker(data, labels, mode='cg', multichannel=True)
    assert data[..., 0].shape == labels.shape
    single_labels = random_walker(data[..., 0], labels, mode='cg')
    assert (multi_labels.reshape(labels.shape)[25:45, 40:60] == 2).all()
    assert data[..., 0].shape == labels.shape
    return data, multi_labels, single_labels, labels

def test_multispectral_3d():
    n = 30
    lx, ly, lz = n, n, n
    data, labels = make_3d_syntheticdata(lx, ly, lz)
    data.shape += (1,)
    data = data.repeat(2, axis=3)   # Result should be identical
    multi_labels = random_walker(data, labels, mode='cg', multichannel=True)
    single_labels = random_walker(data[..., 0], labels, mode='cg')
    assert (multi_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all()
    assert (single_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all()
    return data, multi_labels, single_labels, labels

def test_multispectral_3d():
    n = 30
    lx, ly, lz = n, n, n
    data, labels = make_3d_syntheticdata(lx, ly, lz)
    data = data[..., np.newaxis].repeat(2, axis=-1)  # Expect identical output
    multi_labels = random_walker(data, labels, mode='cg', multichannel=True)
    assert data[..., 0].shape == labels.shape
    single_labels = random_walker(data[..., 0], labels, mode='cg')
    assert (multi_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all()
    assert (single_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all()
    assert data[..., 0].shape == labels.shape
    return data, multi_labels, single_labels, labels

def test_spacing_1():
    n = 30
    lx, ly, lz = n, n, n
    data, _ = make_3d_syntheticdata(lx, ly, lz)

    # Rescale `data` along Y axis
    # `resize` is not yet 3D capable, so this must be done by looping in 2D.
    data_aniso = np.zeros((n, n * 2, n))
    for i, yz in enumerate(data):
        data_aniso[i, :, :] = resize(yz, (n * 2, n),
                                     mode='constant',
                                     anti_aliasing=False)

    # Generate new labels
    small_l = int(lx // 5)
    labels_aniso = np.zeros_like(data_aniso)
    labels_aniso[lx // 5, ly // 5, lz // 5] = 1
    labels_aniso[lx // 2 + small_l // 4,
                 ly - small_l // 2,
                 lz // 2 - small_l // 4] = 2

    # Test with `spacing` kwarg
    # First, anisotropic along Y
    with expected_warnings(['"cg" mode' + '|' + SCIPY_RANK_WARNING,
                            NUMPY_MATRIX_WARNING]):
        labels_aniso = random_walker(data_aniso, labels_aniso, mode='cg',
                                     spacing=(1., 2., 1.))
    assert (labels_aniso[13:17, 26:34, 13:17] == 2).all()

    # Rescale `data` along X axis
    # `resize` is not yet 3D capable, so this must be done by looping in 2D.
    data_aniso = np.zeros((n, n * 2, n))
    for i in range(data.shape[1]):
        data_aniso[i, :, :] = resize(data[:, 1, :], (n * 2, n),
                                     mode='constant',
                                     anti_aliasing=False)

    # Generate new labels
    small_l = int(lx // 5)
    labels_aniso2 = np.zeros_like(data_aniso)
    labels_aniso2[lx // 5, ly // 5, lz // 5] = 1
    labels_aniso2[lx - small_l // 2,
                  ly // 2 + small_l // 4,
                  lz // 2 - small_l // 4] = 2

    # Anisotropic along X
    with expected_warnings(['"cg" mode' + '|' + SCIPY_RANK_WARNING,
                            NUMPY_MATRIX_WARNING]):
        labels_aniso2 = random_walker(data_aniso,
                                      labels_aniso2,
                                      mode='cg', spacing=(2., 1., 1.))
    assert (labels_aniso2[26:34, 13:17, 13:17] == 2).all()

def test_2d_cg():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels_cg = random_walker(data, labels, beta=90, mode='cg')
    assert (labels_cg[25:45, 40:60] == 2).all()
    assert data.shape == labels.shape
    full_prob = random_walker(data, labels, beta=90, mode='cg',
                              return_full_prob=True)
    assert (full_prob[1, 25:45, 40:60] >=
            full_prob[0, 25:45, 40:60]).all()
    assert data.shape == labels.shape
    return data, labels_cg

def test_trivial_cases():
    # When all voxels are labeled
    img = np.ones((10, 10))
    labels = np.ones((10, 10))
    pass_through = random_walker(img, labels)
    np.testing.assert_array_equal(pass_through, labels)

    # When all voxels are labeled AND return_full_prob is True
    labels[:, :5] = 3
    expected = np.concatenate(((labels == 1)[..., np.newaxis],
                               (labels == 3)[..., np.newaxis]), axis=2)
    test = random_walker(img, labels, return_full_prob=True)
    np.testing.assert_array_equal(test, expected)

def test_multispectral_2d():
    lx, ly = 70, 100
    data, labels = make_2d_syntheticdata(lx, ly)
    data = data[..., np.newaxis].repeat(2, axis=-1)  # Expect identical output
    with expected_warnings(['"cg" mode' + '|' + SCIPY_RANK_WARNING,
                            NUMPY_MATRIX_WARNING]):
        multi_labels = random_walker(data, labels, mode='cg',
                                     multichannel=True)
    assert data[..., 0].shape == labels.shape
    with expected_warnings(['"cg" mode' + '|' + SCIPY_RANK_WARNING,
                            NUMPY_MATRIX_WARNING]):
        single_labels = random_walker(data[..., 0], labels, mode='cg')
    assert (multi_labels.reshape(labels.shape)[25:45, 40:60] == 2).all()
    assert data[..., 0].shape == labels.shape
    return data, multi_labels, single_labels, labels

def test_2d_cg():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    with expected_warnings(['"cg" mode' + '|' + SCIPY_EXPECTED]):
        labels_cg = random_walker(data, labels, beta=90, mode='cg')
    assert (labels_cg[25:45, 40:60] == 2).all()
    assert data.shape == labels.shape
    with expected_warnings(['"cg" mode' + '|' + SCIPY_EXPECTED]):
        full_prob = random_walker(data, labels, beta=90, mode='cg',
                                  return_full_prob=True)
    assert (full_prob[1, 25:45, 40:60] >=
            full_prob[0, 25:45, 40:60]).all()
    assert data.shape == labels.shape
    return data, labels_cg

def test_3d():
    n = 30
    lx, ly, lz = n, n, n
    data, labels = make_3d_syntheticdata(lx, ly, lz)
    labels = random_walker(data, labels, mode='cg')
    assert (labels.reshape(data.shape)[13:17, 13:17, 13:17] == 2).all()
    return data, labels

def test_2d_cg_mg():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels_cg_mg = random_walker(data, labels, beta=90, mode='cg_mg')
    assert (labels_cg_mg[25:45, 40:60] == 2).all()
    return data, labels_cg_mg

def cluster_by_diffusion(data):
    markers = np.zeros(data.shape, dtype=np.uint)
    markers[data < -0.00] = 1
    markers[data > 0.03] = 2
    labels2 = random_walker(data, markers, beta=10, mode='bf')

    return labels2

def start():
    # Generate noisy synthetic data
    data = microstructure(l=128)
    data += 0.35 * np.random.randn(*data.shape)
    markers = np.zeros(data.shape, dtype=np.uint)
    markers[data < -0.3] = 1
    markers[data > 1.3] = 2

    # Run random walker algorithm
    labels = random_walker(data, markers, beta=10, mode='bf')

    # Plot results
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(8, 3.2))
    ax1.imshow(data, cmap='gray', interpolation='nearest')
    ax1.axis('off')
    ax1.set_title('Noisy data')
    ax2.imshow(markers, cmap='hot', interpolation='nearest')
    ax2.axis('off')
    ax2.set_title('Markers')
    ax3.imshow(labels, cmap='gray', interpolation='nearest')
    ax3.axis('off')
    ax3.set_title('Segmentation')

    fig.subplots_adjust(hspace=0.01, wspace=0.01, top=1, bottom=0, left=0,
                        right=1)
    plt.show()

def segment_out_cells(base):
    # TODO: try using OTSU for GFP thresholding

    sel_elem = disk(2)
    gfp_collector = np.sum(base, axis=0)
    gfp_clustering_markers = np.zeros(gfp_collector.shape, dtype=np.uint8)
    # random walker segment
    gfp_clustering_markers[gfp_collector > np.mean(gfp_collector) * 2] = 2
    gfp_clustering_markers[gfp_collector < np.mean(gfp_collector) * 0.20] = 1
    labels = random_walker(gfp_collector, gfp_clustering_markers, beta=10, mode='bf')
    # round up the labels and set the background to 0 from 1
    labels = closing(labels, sel_elem)
    labels -= 1
    # prepare distances for the watershed
    distance = ndi.distance_transform_edt(labels)
    local_maxi = peak_local_max(distance,
                                indices=False,  # we want the image mask, not peak position
                                min_distance=10,  # about half of a bud with our size
                                threshold_abs=10,  # allows to clear the noise
                                labels=labels)
    # we fuse the labels that are close together that escaped the min distance in local_maxi
    local_maxi = ndi.convolve(local_maxi, np.ones((5, 5)), mode='constant', cval=0.0)
    # finish the watershed
    expanded_maxi_markers = ndi.label(local_maxi, structure=np.ones((3, 3)))[0]
    segmented_cells_labels = watershed(-distance, expanded_maxi_markers, mask=labels)

    # log debugging data
    running_debug_frame.gfp_collector = gfp_collector
    running_debug_frame.gfp_clustering_markers = gfp_clustering_markers
    running_debug_frame.labels = labels
    running_debug_frame.segmented_cells_labels = segmented_cells_labels

    return gfp_collector, segmented_cells_labels

    def _automatic_localization(self):
        """
        Automatic localization made by Tomas Ryba.
        """

        # seeds = self.get_seeds_using_class_1(class1)
        liver = self.data3d * (self.segmentation != 0)
        print('analyzing histogram...')
        class1 = tools.analyse_histogram(self.data3d,
                                         roi=self.segmentation != 0)
        # sed3.sed3(self.data3d, seeds=class1).show()
        print('getting seeds...')
        seeds = self.get_seeds_using_prob_class1(
            liver,
            class1,
            thresholdType='percOfMaxDist',
            percT=0.3)

        # sed3.sed3(self.data3d, seeds=seeds).show()

        print('Starting random walker...')
        rw = random_walker(liver, seeds, mode='cg_mg')
        print('...finished.')

        label_l = self.data['slab']['lesions']

        lessions = rw == 2
        sed3.sed3(self.data3d, contour=lessions).show()
        lessions = self.filter_objects(lessions)

        self.segmentation = np.where(lessions, label_l, self.segmentation)

def test_2d_bf():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels_bf = random_walker(data, labels, beta=90, mode="bf")
    assert (labels_bf[25:45, 40:60] == 2).all()
    assert data.shape == labels.shape
    full_prob_bf = random_walker(data, labels, beta=90, mode="bf", return_full_prob=True)
    assert (full_prob_bf[1, 25:45, 40:60] >= full_prob_bf[0, 25:45, 40:60]).all()
    assert data.shape == labels.shape
    # Now test with more than two labels
    labels[55, 80] = 3
    full_prob_bf = random_walker(data, labels, beta=90, mode="bf", return_full_prob=True)
    assert (full_prob_bf[1, 25:45, 40:60] >= full_prob_bf[0, 25:45, 40:60]).all()
    assert len(full_prob_bf) == 3
    assert data.shape == labels.shape

def test_spacing_0():
    n = 30
    lx, ly, lz = n, n, n
    data, _ = make_3d_syntheticdata(lx, ly, lz)

    # Rescale `data` along Z axis
    data_aniso = np.zeros((n, n, n // 2))
    for i, yz in enumerate(data):
        data_aniso[i, :, :] = resize(yz, (n, n // 2),
                                     mode='constant',
                                     anti_aliasing=False)

    # Generate new labels
    small_l = int(lx // 5)
    labels_aniso = np.zeros_like(data_aniso)
    labels_aniso[lx // 5, ly // 5, lz // 5] = 1
    labels_aniso[lx // 2 + small_l // 4,
                 ly // 2 - small_l // 4,
                 lz // 4 - small_l // 8] = 2

    # Test with `spacing` kwarg
    with expected_warnings(['"cg" mode' + '|' + SCIPY_RANK_WARNING,
                            NUMPY_MATRIX_WARNING]):
        labels_aniso = random_walker(data_aniso, labels_aniso, mode='cg',
                                     spacing=(1., 1., 0.5))

    assert (labels_aniso[13:17, 13:17, 7:9] == 2).all()