def test_edge_feature_latent_node_crf_no_latent():
    # no latent nodes

    # Test inference with different weights in different directions

    X, Y = toy.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))
        assert_array_almost_equal(res[1], y_pred[1])
        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_multinomial_blocks():
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.3, seed=0)
    crf = GridCRF(n_states=X.shape[-1])
    clf = StructuredPerceptron(problem=crf, max_iter=10)
    clf.fit(X, Y)
    Y_pred = clf.predict(X)
    assert_array_equal(Y, Y_pred)

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_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_switch_to_ad3():
    # test if switching between qpbo and ad3 works

    if not get_installed(['qpbo']) or not get_installed(['ad3']):
        return
    X, Y = toy.generate_blocks_multinomial(n_samples=5, noise=1.5,
                                           seed=0)
    crf = GridCRF(n_states=3, inference_method='qpbo')

    ssvm = NSlackSSVM(crf, max_iter=10000)

    ssvm_with_switch = NSlackSSVM(crf, max_iter=10000, switch_to=('ad3'))
    ssvm.fit(X, Y)
    ssvm_with_switch.fit(X, Y)
    assert_equal(ssvm_with_switch.model.inference_method, 'ad3')
    # we check that the dual is higher with ad3 inference
    # as it might use the relaxation, that is pretty much guraranteed
    assert_greater(ssvm_with_switch.objective_curve_[-1],
                   ssvm.objective_curve_[-1])
    print(ssvm_with_switch.objective_curve_[-1], ssvm.objective_curve_[-1])

    # test that convergence also results in switch
    ssvm_with_switch = NSlackSSVM(crf, max_iter=10000, switch_to=('ad3'),
                                  tol=10)
    ssvm_with_switch.fit(X, Y)
    assert_equal(ssvm_with_switch.model.inference_method, 'ad3')

def test_blocks_multinomial_crf():
    X, Y = toy.generate_blocks_multinomial(n_samples=1)
    x, y = X[0], Y[0]
    w = np.array([1.0, 1.0, 1.0, 0.4, -0.3, 0.3, -0.5, -0.1, 0.3])
    crf = GridCRF(n_states=3)
    y_hat = crf.inference(x, w)
    assert_array_equal(y, y_hat)

def test_psi_continuous():
    # 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]

    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
    crf = DirectionalGridCRF(n_states=3, inference_method='lp')
    w = np.hstack([np.ones(3), -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,))
    # first unary, then horizontal, then vertical
    unary_psi = crf.get_unary_weights(psi_y)
    pw_psi_horz, pw_psi_vert = crf.get_pairwise_weights(psi_y)

    # test unary
    xx, yy = np.indices(y.shape)
    assert_array_almost_equal(unary_psi,
                              np.bincount(y.ravel(), x[xx, yy, y].ravel()))

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_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]

    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 get_installed(["lp", "ad3"]):
        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)

        # 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]

    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(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)
        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)
        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_multinomial_blocks_cutting_plane():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.3,
                                           seed=0)
    n_labels = len(np.unique(Y))
    crf = GridCRF(n_states=n_labels)
    clf = StructuredSVM(problem=crf, max_iter=10, C=100, verbose=0,
                        check_constraints=False)
    clf.fit(X, Y)
    Y_pred = clf.predict(X)
    assert_array_equal(Y, Y_pred)

def test_multinomial_blocks_directional():
    # 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)
    n_labels = len(np.unique(Y))
    crf = DirectionalGridCRF(n_states=n_labels)
    clf = NSlackSSVM(model=crf, max_iter=100, C=100, verbose=0,
                     check_constraints=True, batch_size=1)
    clf.fit(X, Y)
    Y_pred = clf.predict(X)
    assert_array_equal(Y, Y_pred)

def test_multinomial_blocks_subgradient():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.3,
                                           seed=1)
    n_labels = len(np.unique(Y))
    crf = GridCRF(n_states=n_labels)
    clf = SubgradientSSVM(model=crf, max_iter=50, C=10, momentum=.98,
                          learning_rate=0.001)
    clf.fit(X, Y)
    Y_pred = clf.predict(X)
    assert_array_equal(Y, Y_pred)

def test_multinomial_blocks_cutting_plane():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=40, noise=0.5, seed=0)
    n_labels = len(np.unique(Y))
    for inference_method in get_installed(['ad3']):
        crf = GridCRF(n_states=n_labels, inference_method=inference_method)
        clf = NSlackSSVM(model=crf, max_iter=100, C=100, verbose=0,
                         check_constraints=False, batch_size=1)
        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        assert_array_equal(Y, Y_pred)

def test_multinomial_blocks_one_slack():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.3,
                                           seed=0)
    n_labels = len(np.unique(Y))
    for inference_method in ['lp']:
        crf = GridCRF(n_states=n_labels, inference_method=inference_method)
        clf = OneSlackSSVM(problem=crf, max_iter=50, C=100, verbose=100,
                           check_constraints=True, break_on_bad=True)
        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        assert_array_equal(Y, Y_pred)

def test_multinomial_blocks_cutting_plane():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.3,
                                           seed=0)
    n_labels = len(np.unique(Y))
    for inference_method in ['lp', 'qpbo', 'ad3']:
        crf = GridCRF(n_states=n_labels, inference_method=inference_method)
        clf = StructuredSVM(model=crf, max_iter=10, C=100, verbose=0,
                            check_constraints=False, n_jobs=-1)
        clf.fit(X, Y)
        Y_pred = clf.predict(X)
        assert_array_equal(Y, Y_pred)

def test_blocks_multinomial_crf():
    X, Y = toy.generate_blocks_multinomial(n_samples=1, size_x=9, seed=0)
    x, y = X[0], Y[0]
    w = np.array([1., 0., 0.,  # unaryA
                  0., 1., 0.,
                  0., 0., 1.,
                 .4,           # pairwise
                 -.3, .3,
                 -.5, -.1, .3])
    for inference_method in get_installed():
        crf = GridCRF(n_states=3, inference_method=inference_method)
        y_hat = crf.inference(x, w)
        assert_array_equal(y, y_hat)

def test_averaged():
    # Under a lot of noise, averaging helps.  This fails with less noise.
    X, Y = toy.generate_blocks_multinomial(n_samples=15, noise=2, seed=0)
    X_train, Y_train = X[:10], Y[:10]
    X_test, Y_test = X[10:], Y[10:]
    crf = GridCRF(n_states=X.shape[-1])
    clf = StructuredPerceptron(model=crf, max_iter=3)
    clf.fit(X_train, Y_train)
    no_avg_test = clf.score(X_test, Y_test)
    clf.set_params(average=True)
    clf.fit(X_train, Y_train)
    avg_test = clf.score(X_test, Y_test)
    assert_greater(avg_test, no_avg_test)

def test_multinomial_blocks_one_slack():
    #testing cutting plane ssvm on easy multinomial dataset
    X, Y = toy.generate_blocks_multinomial(n_samples=10, noise=0.5,
                                           seed=0)
    print(np.argmax(X[0], axis=-1))
    n_labels = len(np.unique(Y))
    crf = GridCRF(n_states=n_labels)
    clf = OneSlackSSVM(model=crf, max_iter=150, C=1,
                       check_constraints=True, break_on_bad=True, tol=.1,
                       inference_cache=50)
    clf.fit(X, Y)
    Y_pred = clf.predict(X)
    assert_array_equal(Y, Y_pred)

def test_blocks_multinomial_crf():
    X, Y = toy.generate_blocks_multinomial(n_samples=1)
    x, y = X[0], Y[0]
    w = np.array([1., 0., 0.,  # unaryA
                  0., 1., 0.,
                  0., 0., 1.,
                 .4,           # pairwise
                 -.3, .3,
                 -.5, -.1, .3])
    for inference_method in ['dai', 'qpbo', 'lp', 'ad3']:
        crf = GridCRF(n_states=3, inference_method=inference_method)
        y_hat = crf.inference(x, w)
        assert_array_equal(y, y_hat)