import numpy as np


def test_rel_stoch_order_test():
    from soe.mvtesting import MVStochasticOrderTesting

    k = 3  # number of models
    n_resp = 5
    n_samples = 100  # number of scores per model

    # Generate random scores with a given random mean
    means = np.random.permutation(k)
    gt_rank = np.argsort(means)[::-1]

    # For every model we have its scores in a list
    scores_list = [m + np.random.randn(n_samples, n_resp) for m in means]

    # Compute relative test
    tst = MVStochasticOrderTesting(scores_list,
                                   n_bootstrap=100,
                                   use_sinkhorn=True,
                                   cost='logistic',
                                   verbose=True)

    # Relative stochastic order test
    rank_rel = tst.compute_relative_test(alpha=0.05)

    # Check rank obtained with relative test quantiles
    assert all(rank_rel == gt_rank)


def test_rel_stoch_order_test_cuda():
    from soe.mvtesting import MVStochasticOrderTesting

    k = 3  # number of models
    n_resp = 5
    n_samples = 100  # number of scores per model

    # Generate random scores with a given random mean
    means = np.random.permutation(k)
    gt_rank = np.argsort(means)[::-1]

    # For every model we have its scores in a list
    scores_list = [m + np.random.randn(n_samples, n_resp) for m in means]

    # Compute relative test
    tst = MVStochasticOrderTesting(scores_list,
                                   n_bootstrap=100,
                                   use_sinkhorn=True,
                                   use_cuda=True,
                                   cost='logistic',
                                   verbose=True)

    # Relative stochastic order test
    rank_rel = tst.compute_relative_test(alpha=0.05)

    # Check rank obtained with relative test quantiles
    assert all(rank_rel == gt_rank)
