def test_latent_node_boxes_edge_features():
    # learn the "easy" 2x2 boxes dataset.
    # smoketest using a single constant edge feature

    X, Y = make_simple_2x2(seed=1, n_samples=40)
    latent_crf = EdgeFeatureLatentNodeCRF(n_labels=2, n_hidden_states=2, n_features=1)
    base_svm = OneSlackSSVM(latent_crf)
    base_svm.C = 10
    latent_svm = LatentSSVM(base_svm,
                            latent_iter=10)

    G = [make_grid_edges(x) for x in X]

    # make edges for hidden states:
    edges = make_edges_2x2()

    G = [np.vstack([make_grid_edges(x), edges]) for x in X]

    # reshape / flatten x and y
    X_flat = [x.reshape(-1, 1) for x in X]
    Y_flat = [y.ravel() for y in Y]

    #X_ = zip(X_flat, G, [2 * 2 for x in X_flat])
    # add edge features
    X_ = [(x, g, np.ones((len(g), 1)), 4) for x, g in zip(X_flat, G)]
    latent_svm.fit(X_[:20], Y_flat[:20])

    assert_array_equal(latent_svm.predict(X_[:20]), Y_flat[:20])
    assert_equal(latent_svm.score(X_[:20], Y_flat[:20]), 1)

    # test that score is not always 1
    assert_true(.98 < latent_svm.score(X_[20:], Y_flat[20:]) < 1)

def test_latent_node_boxes_standard_latent():
    # learn the "easy" 2x2 boxes dataset.
    # a 2x2 box is placed randomly in a 4x4 grid
    # we add a latent variable for each 2x2 patch
    # that should make the model fairly simple

    X, Y = make_simple_2x2(seed=1, n_samples=40)
    latent_crf = LatentNodeCRF(n_labels=2, n_hidden_states=2, n_features=1)
    one_slack = OneSlackSSVM(latent_crf)
    n_slack = NSlackSSVM(latent_crf)
    subgradient = SubgradientSSVM(latent_crf, max_iter=100)
    for base_svm in [one_slack, n_slack, subgradient]:
        base_svm.C = 10
        latent_svm = LatentSSVM(base_svm,
                                latent_iter=10)

        G = [make_grid_edges(x) for x in X]

        # make edges for hidden states:
        edges = make_edges_2x2()

        G = [np.vstack([make_grid_edges(x), edges]) for x in X]

        # reshape / flatten x and y
        X_flat = [x.reshape(-1, 1) for x in X]
        Y_flat = [y.ravel() for y in Y]

        X_ = list(zip(X_flat, G, [2 * 2 for x in X_flat]))
        latent_svm.fit(X_[:20], Y_flat[:20])

        assert_array_equal(latent_svm.predict(X_[:20]), Y_flat[:20])
        assert_equal(latent_svm.score(X_[:20], Y_flat[:20]), 1)

        # test that score is not always 1
        assert_true(.98 < latent_svm.score(X_[20:], Y_flat[20:]) < 1)

def test_latent_node_boxes_latent_subgradient():
    # same as above, now with elementary subgradients

    # learn the "easy" 2x2 boxes dataset.
    # a 2x2 box is placed randomly in a 4x4 grid
    # we add a latent variable for each 2x2 patch
    # that should make the model fairly simple

    X, Y = toy.make_simple_2x2(seed=1)
    latent_crf = LatentNodeCRF(n_labels=2, inference_method='lp',
                               n_hidden_states=2, n_features=1)
    latent_svm = LatentSubgradientSSVM(model=latent_crf, max_iter=250, C=10,
                                       verbose=10, learning_rate=0.1,
                                       momentum=0)

    G = [make_grid_edges(x) for x in X]

    # make edges for hidden states:
    edges = []
    node_indices = np.arange(4 * 4).reshape(4, 4)
    for i, (x, y) in enumerate(itertools.product([0, 2], repeat=2)):
        for j in xrange(x, x + 2):
            for k in xrange(y, y + 2):
                edges.append([i + 4 * 4, node_indices[j, k]])

    G = [np.vstack([make_grid_edges(x), edges]) for x in X]

    # reshape / flatten x and y
    X_flat = [x.reshape(-1, 1) for x in X]
    Y_flat = [y.ravel() for y in Y]

    X_ = zip(X_flat, G, [4 * 4 for x in X_flat])
    latent_svm.fit(X_, Y_flat)

    assert_equal(latent_svm.score(X_, Y_flat), 1)

def test_joint_feature_discrete():
    """
    Testing with a single type of nodes. Must de aw well as EdgeFeatureGraphCRF
    """
    X, Y = generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)
    edge_features = edge_list_to_features(edge_list)
    x = ([x.reshape(-1, 3)], [edges], [edge_features])
    y_flat = y.ravel()
    #for inference_method in get_installed(["lp", "ad3", "qpbo"]):
    if True:
        crf = NodeTypeEdgeFeatureGraphCRF(1, [3], [3], [[2]])
        joint_feature_y = crf.joint_feature(x, y_flat)
        assert_equal(joint_feature_y.shape, (crf.size_joint_feature,))
        # first horizontal, then vertical
        # we trust the unaries ;)
        n_states = crf.l_n_states[0]
        n_features = crf.l_n_features[0]
        pw_joint_feature_horz, pw_joint_feature_vert = joint_feature_y[n_states *
                                         n_features:].reshape(
                                             2, n_states, n_states)
        assert_array_equal(pw_joint_feature_vert, np.diag([9 * 4, 9 * 4, 9 * 4]))
        vert_joint_feature = np.diag([10 * 3, 10 * 3, 10 * 3])
        vert_joint_feature[0, 1] = 10
        vert_joint_feature[1, 2] = 10
        assert_array_equal(pw_joint_feature_horz, vert_joint_feature)

def test_binary_blocks_cutting_plane():
    #testing cutting plane ssvm on easy binary dataset
    # generate graphs explicitly for each example
    for inference_method in get_installed(["lp", "qpbo", "ad3", 'ogm']):
        X, Y = generate_blocks(n_samples=3)
        crf = GraphCRF(inference_method=inference_method)
        clf = NSlackSSVM(model=crf, max_iter=20, C=100, check_constraints=True,
                         break_on_bad=False, n_jobs=1)
        x1, x2, x3 = X
        y1, y2, y3 = Y
        n_states = len(np.unique(Y))
        # delete some rows to make it more fun
        x1, y1 = x1[:, :-1], y1[:, :-1]
        x2, y2 = x2[:-1], y2[:-1]
        # generate graphs
        X_ = [x1, x2, x3]
        G = [make_grid_edges(x) for x in X_]

        # reshape / flatten x and y
        X_ = [x.reshape(-1, n_states) for x in X_]
        Y = [y.ravel() for y in [y1, y2, y3]]

        X = list(zip(X_, G))

        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        for y, y_pred in zip(Y, Y_pred):
            assert_array_equal(y, y_pred)

def prepare_data(X):
    X_directions = []
    X_edge_features = []
    for x in X:
        # get edges in grid
        right, down = make_grid_edges(x, return_lists=True)
        edges = np.vstack([right, down])
        # use 3x3 patch around each point
        features = neighborhood_feature(x)
        # simple edge feature that encodes just if an edge is horizontal or
        # vertical
        edge_features_directions = edge_list_to_features([right, down])
        # edge feature that contains features from the nodes that the edge connects
        edge_features = np.zeros((edges.shape[0], features.shape[1], 4))
        edge_features[:len(right), :, 0] = features[right[:, 0]]
        edge_features[:len(right), :, 1] = features[right[:, 1]]
#---ORIGINAL CODE        
#         edge_features[len(right):, :, 0] = features[down[:, 0]]
#         edge_features[len(right):, :, 1] = features[down[:, 1]]
        edge_features[len(right):, :, 2] = features[down[:, 0]]
        edge_features[len(right):, :, 3] = features[down[:, 1]]
#---END OF FIX        
        edge_features = edge_features.reshape(edges.shape[0], -1)
        X_directions.append((features, edges, edge_features_directions))
        X_edge_features.append((features, edges, edge_features))
    return X_directions, X_edge_features

def test_energy_discrete():
#     for inference_method in get_installed(["qpbo", "ad3"]):
#         crf = EdgeFeatureGraphCRF(n_states=3,
#                                   inference_method=inference_method,
#                                   n_edge_features=2, n_features=3)
        crf = NodeTypeEdgeFeatureGraphCRF(1, [3], [3], [[2]])
        
        for i in range(10):
            x = np.random.normal(size=(7, 8, 3))
            edge_list = make_grid_edges(x, 4, return_lists=True)
            edges = np.vstack(edge_list)
            edge_features = edge_list_to_features(edge_list)
            x = ([x.reshape(-1, 3)], [edges], [edge_features])

            unary_params = np.random.normal(size=(3, 3))
            pw1 = np.random.normal(size=(3, 3))
            pw2 = np.random.normal(size=(3, 3))
            w = np.hstack([unary_params.ravel(), pw1.ravel(), pw2.ravel()])
            crf.initialize(x)
            y_hat = crf.inference(x, w, relaxed=False)
            #flat_edges = crf._index_all_edges(x)
            energy = compute_energy(crf._get_unary_potentials(x, w)[0],
                                    crf._get_pairwise_potentials(x, w)[0], edges, #CAUTION: pass the flatened edges!!
                                    y_hat)

            joint_feature = crf.joint_feature(x, y_hat)
            energy_svm = np.dot(joint_feature, w)

            assert_almost_equal(energy, energy_svm)

def test_energy_continuous():
    # make sure that energy as computed by ssvm is the same as by lp
    np.random.seed(0)
    for inference_method in get_installed(["lp", "ad3"]):
        found_fractional = False
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2, n_features=3)
        while not found_fractional:
            x = np.random.normal(size=(7, 8, 3))
            edge_list = make_grid_edges(x, 4, return_lists=True)
            edges = np.vstack(edge_list)
            edge_features = edge_list_to_features(edge_list)
            x = (x.reshape(-1, 3), edges, edge_features)

            unary_params = np.random.normal(size=(3, 3))
            pw1 = np.random.normal(size=(3, 3))
            pw2 = np.random.normal(size=(3, 3))
            w = np.hstack([unary_params.ravel(), pw1.ravel(), pw2.ravel()])
            res, energy = crf.inference(x, w, relaxed=True, return_energy=True)
            found_fractional = np.any(np.max(res[0], axis=-1) != 1)

            joint_feature = crf.joint_feature(x, res)
            energy_svm = np.dot(joint_feature, w)

            assert_almost_equal(energy, -energy_svm)

def test_binary_blocks_one_slack_graph():
    #testing cutting plane ssvm on easy binary dataset
    # generate graphs explicitly for each example
    for inference_method in ["dai", "lp"]:
        print("testing %s" % inference_method)
        X, Y = toy.generate_blocks(n_samples=3)
        crf = GraphCRF(inference_method=inference_method)
        clf = OneSlackSSVM(problem=crf, max_iter=100, C=100, verbose=100,
                           check_constraints=True, break_on_bad=True,
                           n_jobs=1)
        x1, x2, x3 = X
        y1, y2, y3 = Y
        n_states = len(np.unique(Y))
        # delete some rows to make it more fun
        x1, y1 = x1[:, :-1], y1[:, :-1]
        x2, y2 = x2[:-1], y2[:-1]
        # generate graphs
        X_ = [x1, x2, x3]
        G = [make_grid_edges(x) for x in X_]

        # reshape / flatten x and y
        X_ = [x.reshape(-1, n_states) for x in X_]
        Y = [y.ravel() for y in [y1, y2, y3]]

        X = zip(X_, G)

        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        for y, y_pred in zip(Y, Y_pred):
            assert_array_equal(y, y_pred)

def test_edge_feature_latent_node_crf_no_latent():
    # no latent nodes

    # Test inference with different weights in different directions

    X, Y = generate_blocks_multinomial(noise=2, n_samples=1, seed=1, size_x=10)
    x, y = X[0], Y[0]
    n_states = x.shape[-1]

    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)

    pw_horz = -1 * np.eye(n_states + 5)
    xx, yy = np.indices(pw_horz.shape)
    # linear ordering constraint horizontally
    pw_horz[xx > yy] = 1

    # high cost for unequal labels vertically
    pw_vert = -1 * np.eye(n_states + 5)
    pw_vert[xx != yy] = 1
    pw_vert *= 10

    # generate edge weights
    edge_weights_horizontal = np.repeat(pw_horz[np.newaxis, :, :],
                                        edge_list[0].shape[0], axis=0)
    edge_weights_vertical = np.repeat(pw_vert[np.newaxis, :, :],
                                      edge_list[1].shape[0], axis=0)
    edge_weights = np.vstack([edge_weights_horizontal, edge_weights_vertical])

    # do inference
    # pad x for hidden states...
    x_padded = -100 * np.ones((x.shape[0], x.shape[1], x.shape[2] + 5))
    x_padded[:, :, :x.shape[2]] = x
    res = lp_general_graph(-x_padded.reshape(-1, n_states + 5), edges,
                           edge_weights)

    edge_features = edge_list_to_features(edge_list)
    x = (x.reshape(-1, n_states), edges, edge_features, 0)
    y = y.ravel()

    for inference_method in get_installed(["lp"]):
        # same inference through CRF inferface
        crf = EdgeFeatureLatentNodeCRF(n_labels=3,
                                       inference_method=inference_method,
                                       n_edge_features=2, n_hidden_states=5)
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=True)
        assert_array_almost_equal(res[0], y_pred[0].reshape(-1, n_states + 5),
                                  4)
        assert_array_almost_equal(res[1], y_pred[1], 4)
        assert_array_equal(y, np.argmax(y_pred[0], axis=-1))

    for inference_method in get_installed(["lp", "ad3", "qpbo"]):
        # again, this time discrete predictions only
        crf = EdgeFeatureLatentNodeCRF(n_labels=3,
                                       inference_method=inference_method,
                                       n_edge_features=2, n_hidden_states=5)
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=False)
        assert_array_equal(y, y_pred)

def test_energy_discrete():
    for inference_method in get_installed(["qpbo", "ad3"]):
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2)
        for i in xrange(10):
            x = np.random.normal(size=(7, 8, 3))
            edge_list = make_grid_edges(x, 4, return_lists=True)
            edges = np.vstack(edge_list)
            edge_features = edge_list_to_features(edge_list)
            x = (x.reshape(-1, 3), edges, edge_features)

            unary_params = np.random.normal(size=(3, 3))
            pw1 = np.random.normal(size=(3, 3))
            pw2 = np.random.normal(size=(3, 3))
            w = np.hstack([unary_params.ravel(), pw1.ravel(), pw2.ravel()])
            y_hat = crf.inference(x, w, relaxed=False)
            energy = compute_energy(crf.get_unary_potentials(x, w),
                                    crf.get_pairwise_potentials(x, w), edges,
                                    y_hat)

            psi = crf.psi(x, y_hat)
            energy_svm = np.dot(psi, w)

            assert_almost_equal(energy, energy_svm)

def test_energy():
    # make sure that energy as computed by ssvm is the same as by lp
    np.random.seed(0)
    for inference_method in ["lp", "ad3"]:
        found_fractional = False
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2)
        while not found_fractional:
            x = np.random.normal(size=(7, 8, 3))
            edge_list = make_grid_edges(x, 4, return_lists=True)
            edges = np.vstack(edge_list)
            edge_features = edge_list_to_features(edge_list)
            x = (x.reshape(-1, 3), edges, edge_features)

            unary_params = np.random.normal(size=(3, 3))
            pw1 = np.random.normal(size=(3, 3))
            pw2 = np.random.normal(size=(3, 3))
            w = np.hstack([unary_params.ravel(), pw1.ravel(), pw2.ravel()])
            res, energy = crf.inference(x, w, relaxed=True, return_energy=True)
            found_fractional = np.any(np.max(res[0], axis=-1) != 1)

            psi = crf.psi(x, res)
            energy_svm = np.dot(psi, w)

            assert_almost_equal(energy, -energy_svm)
            if not found_fractional:
                # exact discrete labels, test non-relaxed version
                res, energy = crf.inference(x, w, relaxed=False,
                                            return_energy=True)
                psi = crf.psi(x, res)
                energy_svm = np.dot(psi, w)

                assert_almost_equal(energy, -energy_svm)

def test_psi_discrete():
    X, Y = toy.generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)
    edge_features = edge_list_to_features(edge_list)
    x = (x.reshape(-1, 3), edges, edge_features)
    y_flat = y.ravel()
    for inference_method in ["lp", "ad3", "qpbo"]:
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2)
        psi_y = crf.psi(x, y_flat)
        assert_equal(psi_y.shape, (crf.size_psi,))
        # first horizontal, then vertical
        # we trust the unaries ;)
        pw_psi_horz, pw_psi_vert = psi_y[crf.n_states *
                                         crf.n_features:].reshape(
                                             2, crf.n_states, crf.n_states)
        xx, yy = np.indices(y.shape)
        assert_array_equal(pw_psi_vert, np.diag([9 * 4, 9 * 4, 9 * 4]))
        vert_psi = np.diag([10 * 3, 10 * 3, 10 * 3])
        vert_psi[0, 1] = 10
        vert_psi[1, 2] = 10
        assert_array_equal(pw_psi_horz, vert_psi)

def test_initialization():
    X, Y = generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    n_states = x.shape[-1]

    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)

    edge_features = edge_list_to_features(edge_list)
    x = (x.reshape(-1, n_states), edges, edge_features)
    y = y.ravel()
    crf = EdgeFeatureGraphCRF()
    crf.initialize([x], [y])
    assert_equal(crf.n_edge_features, 2)
    assert_equal(crf.n_features, 3)
    assert_equal(crf.n_states, 3)

    crf = EdgeFeatureGraphCRF(n_states=3,
                              n_features=3,
                              n_edge_features=2)
    # no-op
    crf.initialize([x], [y])

    crf = EdgeFeatureGraphCRF(n_states=4,
                              n_edge_features=2)
    # incompatible
    assert_raises(ValueError, crf.initialize, X=[x], Y=[y])

def test_multinomial_blocks_directional_anti_symmetric():
    # testing cutting plane ssvm with directional CRF on easy multinomial
    # dataset
    X_, Y_ = toy.generate_blocks_multinomial(n_samples=10, noise=0.3, seed=0)
    G = [make_grid_edges(x, return_lists=True) for x in X_]
    edge_features = [edge_list_to_features(edge_list) for edge_list in G]
    edges = [np.vstack(g) for g in G]
    X = zip([x.reshape(-1, 3) for x in X_], edges, edge_features)
    Y = [y.ravel() for y in Y_]

    for inference_method in ['lp', 'ad3']:
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2,
                                  symmetric_edge_features=[0],
                                  antisymmetric_edge_features=[1])
        clf = StructuredSVM(model=crf, max_iter=20, C=1000, verbose=10,
                            check_constraints=False, n_jobs=-1)
        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        assert_array_equal(Y, Y_pred)
        pairwise_params = clf.w[-9 * 2:].reshape(2, 3, 3)
        sym = pairwise_params[0]
        antisym = pairwise_params[1]
        print(sym)
        print(antisym)
        assert_array_equal(sym, sym.T)
        assert_array_equal(antisym, -antisym.T)

def test_psi_continuous():
    # FIXME
    # first make perfect prediction, including pairwise part
    X, Y = toy.generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    n_states = x.shape[-1]
    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)
    edge_features = edge_list_to_features(edge_list)
    x = (x.reshape(-1, 3), edges, edge_features)
    y = y.ravel()

    pw_horz = -1 * np.eye(n_states)
    xx, yy = np.indices(pw_horz.shape)
    # linear ordering constraint horizontally
    pw_horz[xx > yy] = 1

    # high cost for unequal labels vertically
    pw_vert = -1 * np.eye(n_states)
    pw_vert[xx != yy] = 1
    pw_vert *= 10

    # create crf, assemble weight, make prediction
    for inference_method in ["lp", "ad3"]:
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2)
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=True)

        # compute psi for prediction
        psi_y = crf.psi(x, y_pred)
        assert_equal(psi_y.shape, (crf.size_psi,))

def test_inference():
    # Test inference with different weights in different directions

    X, Y = toy.generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    n_states = x.shape[-1]
    edges = make_grid_edges(x, neighborhood=4)

    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)

    pw_horz = -1 * np.eye(n_states)
    xx, yy = np.indices(pw_horz.shape)
    # linear ordering constraint horizontally
    pw_horz[xx > yy] = 1

    # high cost for unequal labels vertically
    pw_vert = -1 * np.eye(n_states)
    pw_vert[xx != yy] = 1
    pw_vert *= 10

    # generate edge weights
    edge_weights_horizontal = np.repeat(pw_horz[np.newaxis, :, :],
                                        edge_list[0].shape[0], axis=0)
    edge_weights_vertical = np.repeat(pw_vert[np.newaxis, :, :],
                                      edge_list[1].shape[0], axis=0)
    edge_weights = np.vstack([edge_weights_horizontal, edge_weights_vertical])

    # do inference
    res = lp_general_graph(-x.reshape(-1, n_states), edges, edge_weights)

    for inference_method in ["lp", "ad3"]:
        # same inference through CRF inferface
        crf = DirectionalGridCRF(n_states=3, inference_method=inference_method)
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=True)
        assert_array_almost_equal(res[0], y_pred[0].reshape(-1, n_states))
        assert_array_almost_equal(res[1], y_pred[1])
        assert_array_equal(y, np.argmax(y_pred[0], axis=-1))

    for inference_method in ["lp", "ad3", "qpbo"]:
        # again, this time discrete predictions only
        crf = DirectionalGridCRF(n_states=3, inference_method=inference_method)
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=False)
        assert_array_equal(y, y_pred)

def test_k_means_initialization_graph_crf():
    # with only 1 state per label, nothing happends
    X, Y = toy.generate_big_checker(n_samples=10)
    crf = LatentGraphCRF(n_labels=2, n_states_per_label=1, inference_method="lp")
    # convert grid model to graph model
    X = [(x.reshape(-1, x.shape[-1]), make_grid_edges(x, return_lists=False)) for x in X]

    H = crf.init_latent(X, Y)
    assert_array_equal(Y, H)

def test_inference():
    # Test inference with different weights in different directions

    X, Y = generate_blocks_multinomial(noise=2, n_samples=1, seed=1)
    x, y = X[0], Y[0]
    n_states = x.shape[-1]

    edge_list = make_grid_edges(x, 4, return_lists=True)
    edges = np.vstack(edge_list)

    pw_horz = -1 * np.eye(n_states)
    xx, yy = np.indices(pw_horz.shape)
    # linear ordering constraint horizontally
    pw_horz[xx > yy] = 1

    # high cost for unequal labels vertically
    pw_vert = -1 * np.eye(n_states)
    pw_vert[xx != yy] = 1
    pw_vert *= 10

    # generate edge weights
    edge_weights_horizontal = np.repeat(pw_horz[np.newaxis, :, :],
                                        edge_list[0].shape[0], axis=0)
    edge_weights_vertical = np.repeat(pw_vert[np.newaxis, :, :],
                                      edge_list[1].shape[0], axis=0)
    edge_weights = np.vstack([edge_weights_horizontal, edge_weights_vertical])

    # do inference
    res = lp_general_graph(-x.reshape(-1, n_states), edges, edge_weights)

    edge_features = edge_list_to_features(edge_list)
    x = (x.reshape(-1, n_states), edges, edge_features)
    y = y.ravel()

    for inference_method in get_installed(["lp", "ad3"]):
        # same inference through CRF inferface
        crf = EdgeFeatureGraphCRF(inference_method=inference_method)
        crf.initialize([x], [y])
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=True)
        if isinstance(y_pred, tuple):
            # ad3 produces an integer result if it found the exact solution
            assert_array_almost_equal(res[1], y_pred[1])
            assert_array_almost_equal(res[0], y_pred[0].reshape(-1, n_states))
            assert_array_equal(y, np.argmax(y_pred[0], axis=-1))

    for inference_method in get_installed(["lp", "ad3", "qpbo"]):
        # again, this time discrete predictions only
        crf = EdgeFeatureGraphCRF(n_states=3,
                                  inference_method=inference_method,
                                  n_edge_features=2)
        crf.initialize([x], [y])
        w = np.hstack([np.eye(3).ravel(), -pw_horz.ravel(), -pw_vert.ravel()])
        y_pred = crf.inference(x, w, relaxed=False)
        assert_array_equal(y, y_pred)

def test_k_means_initialization():
    n_samples = 10
    X, Y = generate_big_checker(n_samples=n_samples)
    edges = [make_grid_edges(x, return_lists=True) for x in X]
    # flatten the grid
    Y = Y.reshape(Y.shape[0], -1)
    X = X.reshape(X.shape[0], -1, X.shape[-1])
    n_labels = len(np.unique(Y))
    X = X.reshape(n_samples, -1, n_labels)

    # sanity check for one state
    H = kmeans_init(X, Y, edges, n_states_per_label=[1] * n_labels,
                    n_labels=n_labels)
    H = np.vstack(H)
    assert_array_equal(Y, H)

    # check number of states
    H = kmeans_init(X, Y, edges, n_states_per_label=[3] * n_labels,
                    n_labels=n_labels)
    H = np.vstack(H)
    assert_array_equal(np.unique(H), np.arange(6))
    assert_array_equal(Y, H // 3)

    # for dataset with more than two states
    X, Y = generate_blocks_multinomial(n_samples=10)
    edges = [make_grid_edges(x, return_lists=True) for x in X]
    Y = Y.reshape(Y.shape[0], -1)
    X = X.reshape(X.shape[0], -1, X.shape[-1])
    n_labels = len(np.unique(Y))

    # sanity check for one state
    H = kmeans_init(X, Y, edges, n_states_per_label=[1] * n_labels,
                    n_labels=n_labels)
    H = np.vstack(H)
    assert_array_equal(Y, H)

    # check number of states
    H = kmeans_init(X, Y, edges, n_states_per_label=[2] * n_labels,
                    n_labels=n_labels)
    H = np.vstack(H)
    assert_array_equal(np.unique(H), np.arange(6))
    assert_array_equal(Y, H // 2)