# Ryan Turner (turnerry@iro.umontreal.ca)
from __future__ import division, print_function

from builtins import range

import numpy as np
import scipy.stats as ss

import mlpaper.constants as cc
import mlpaper.mlpaper as bt
import mlpaper.perf_curves as pc
from mlpaper.classification import DEFAULT_NGRID, curve_boot
from mlpaper.test_constants import FPR
from mlpaper.util import area, interp1d

_FPR = FPR / 3.0  # Divide by number of test funcs


def fail_check_stat(fail, runs, expect_p_fail, fpr):
    pvals_2side = [ss.binom_test(ff, runs, expect_p_fail) for ff in fail]
    pvals_1side = [ss.binom_test(ff, runs, expect_p_fail, alternative="greater") for ff in fail]
    # Note that we are not going multiple comparison correction between the
    # two sided and one sided tests.
    print(fail)
    print(pvals_2side)
    assert np.min(pvals_2side) >= fpr / len(pvals_2side)
    print(pvals_1side)
    assert np.min(pvals_1side) >= fpr / len(pvals_1side)


def test_boot(runs=100):
    N = 201
    confidence = 0.95

    # Drawing more seeds than we need to be safe
    seeds = np.nditer(np.random.randint(low=0, high=int(1e6), size=runs * 5))

    def run_trial(y_true, y_score, y_score_ref, true_curve, curve_f, seed, x_grid=None):
        epsilon = 1e-6

        curve, _ = curve_f(y_true, y_score[:, 1])
        auc, = area(*curve)
        curve, _ = curve_f(y_true, y_score_ref[:, 1])
        auc_ref, = area(*curve)

        true_value, = area(*true_curve)

        np.random.seed(seed)
        (auc_, EB, pval), curve = curve_boot(
            y_true, y_score, ref=true_value, curve_f=curve_f, confidence=confidence, x_grid=x_grid
        )
        true_curve_grid, = interp1d(curve[cc.XGRID].values, *true_curve)
        assert auc_ == auc
        fail_EB = np.abs(auc - true_value) > EB
        # Could also test distn with 1-sided KS test but this easier for now
        fail_P = pval < 1.0 - confidence
        fail_curve = (true_curve_grid < curve[cc.LB].values - epsilon) | (
            curve[cc.UB].values + epsilon < true_curve_grid
        )
        assert (x_grid is None) or np.all(curve[cc.XGRID].values == x_grid)

        np.random.seed(seed)
        (auc_, EB_, pval), curve_ = curve_boot(
            y_true, y_score, ref=y_score_ref, curve_f=curve_f, confidence=confidence, pairwise_CI=False, x_grid=x_grid
        )
        assert auc_ == auc
        assert EB_ == EB
        # Could also test distn with 1-sided KS test but this easier for now
        fail_P2 = pval < 1.0 - confidence
        assert np.all(curve_.values == curve.values)

        np.random.seed(seed)
        (auc_, EB, pval_), curve_ = curve_boot(
            y_true, y_score, ref=y_score_ref, curve_f=curve_f, confidence=confidence, pairwise_CI=True, x_grid=x_grid
        )
        assert auc_ == auc
        fail_EB2 = np.abs(auc - auc_ref) > EB
        # Could also test distn with 1-sided KS test but this easier for now
        assert pval_ == pval
        assert np.all(curve_.values == curve.values)

        return fail_EB, fail_P, fail_EB2, fail_P2, fail_curve

    fail = [0] * 12
    fail_curve_roc = np.zeros(DEFAULT_NGRID, dtype=int)
    fail_curve_ap = np.zeros(DEFAULT_NGRID, dtype=int)
    fail_curve_prg = np.zeros(DEFAULT_NGRID, dtype=int)
    for ii in range(runs):
        mu = np.random.randn(2)
        S = np.random.randn(2, 2)
        S = np.dot(S, S.T)
        # Coverage, esp at edges, is worse for imbalanced data. See issue #20.
        p = 0.5

        x_grid = np.linspace(0.0, 0.99, DEFAULT_NGRID)
        true_curve = (np.array([[0.0, 1.0]]), np.array([[0.0, 1.0]]), pc.LINEAR)
        y_true = np.random.rand(N) <= p
        y_score = np.random.multivariate_normal(mu, S, size=N)
        if np.random.randn() <= 0.5:  # resample to test dupes
            idx = np.random.choice(N, size=N, replace=True)
            y_score = y_score[idx, :]
        y_score, y_score_ref = y_score.T
        y_score = np.stack((np.zeros(N), y_score), axis=1)
        y_score_ref = np.stack((np.zeros(N), y_score_ref), axis=1)

        # Coverage doesn't hold at edges, hence [0.05, 0.95]. See issue #20.
        x_grid = np.linspace(0.05, 0.95, DEFAULT_NGRID)
        fail_EB, fail_P, fail_EB2, fail_P2, fail_curve = run_trial(
            y_true, y_score, y_score_ref, true_curve, pc.roc_curve, next(seeds), x_grid
        )
        fail[0] += fail_EB
        fail[1] += fail_P
        fail[2] += fail_EB2
        fail[3] += fail_P2
        fail_curve_roc += fail_curve

        true_curve = (np.array([[0.0, 1.0]]), np.array([[p, p]]), pc.PREV)
        fail_EB, fail_P, fail_EB2, fail_P2, fail_curve = run_trial(
            y_true, y_score, y_score_ref, true_curve, pc.recall_precision_curve, next(seeds), x_grid
        )
        fail[4] += fail_EB
        fail[5] += fail_P
        fail[6] += fail_EB2
        fail[7] += fail_P2
        fail_curve_ap += fail_curve

        x_grid = np.linspace(0.0, 0.99, DEFAULT_NGRID)
        true_curve = (np.array([[0.0, 1.0]]), np.array([[0.0, 0.0]]), pc.PREV)
        fail_EB, fail_P, fail_EB2, fail_P2, fail_curve = run_trial(
            y_true, y_score, y_score_ref, true_curve, pc.prg_curve, next(seeds), x_grid
        )
        fail[8] += fail_EB
        fail[9] += fail_P
        fail[10] += fail_EB2
        fail[11] += fail_P2
        fail_curve_prg += fail_curve
    sub_FPR = _FPR / 4.0
    expect_p_fail = 1.0 - confidence
    fail_check_stat(fail, runs, expect_p_fail, sub_FPR)
    print("ROC curve")
    fail_check_stat(fail_curve_roc, runs, expect_p_fail, sub_FPR)
    print("RP curve")
    fail_check_stat(fail_curve_ap, runs, expect_p_fail, sub_FPR)
    print("PRG curve")
    fail_check_stat(fail_curve_prg, runs, expect_p_fail, sub_FPR)


def test_boot_mean(runs=100):
    N = 201
    confidence = 0.95

    fail = 0
    for ii in range(runs):
        mu = np.random.randn()
        S = np.abs(np.random.randn())
        x = mu + S * np.random.randn(N)

        mu_est = np.mean(x)
        EB = bt.boot_EB(x, confidence=0.95)

        fail += np.abs(mu - mu_est) > EB
    expect_p_fail = 1.0 - confidence
    print("boot mean")
    fail_check_stat([fail], runs, expect_p_fail, _FPR)


def test_boot_EB_and_test(runs=100):
    """Arguably this should do out to its own file since it tests bt core."""
    mu = np.random.randn()
    stdev = np.abs(np.random.randn())

    N = 201
    confidence = 0.95

    def run_trial(x, true_value):
        _, _, CI = bt._boot_EB_and_test(x, confidence=confidence, return_CI=True)
        LB, UB = CI
        fail_CI = (true_value < LB) or (UB < true_value)

        _, pval, CI = bt._boot_EB_and_test(x - true_value, confidence=confidence, return_CI=True)
        LB, UB = CI
        fail_CI2 = (0 < LB) or (UB < 0)
        fail_P = pval < 1.0 - confidence
        return fail_CI, fail_CI2, fail_P

    fail = [0] * 3
    for ii in range(runs):
        x = mu + stdev * np.random.randn(N)

        fail_CI, fail_CI2, fail_P = run_trial(x, mu)
        fail[0] += fail_CI
        fail[1] += fail_CI2
        fail[2] += fail_P
    expect_p_fail = 1.0 - confidence
    print("boot mean and test")
    fail_check_stat(fail, runs, expect_p_fail, _FPR)


if __name__ == "__main__":
    np.random.seed(56467)

    test_boot()
    test_boot_mean()
    test_boot_EB_and_test()
    print("passed")

[((3118, 3152), 'numpy.zeros', 'np.zeros', (['DEFAULT_NGRID'], {'dtype': 'int'}), '(DEFAULT_NGRID, dtype=int)\n', (3126, 3152), True, 'import numpy as np\n'), ((3173, 3207), 'numpy.zeros', 'np.zeros', (['DEFAULT_NGRID'], {'dtype': 'int'}), '(DEFAULT_NGRID, dtype=int)\n', (3181, 3207), True, 'import numpy as np\n'), ((3229, 3263), 'numpy.zeros', 'np.zeros', (['DEFAULT_NGRID'], {'dtype': 'int'}), '(DEFAULT_NGRID, dtype=int)\n', (3237, 3263), True, 'import numpy as np\n'), ((3278, 3289), 'builtins.range', 'range', (['runs'], {}), '(runs)\n', (3283, 3289), False, 'from builtins import range\n'), ((5816, 5827), 'builtins.range', 'range', (['runs'], {}), '(runs)\n', (5821, 5827), False, 'from builtins import range\n'), ((6294, 6311), 'numpy.random.randn', 'np.random.randn', ([], {}), '()\n', (6309, 6311), True, 'import numpy as np\n'), ((6854, 6865), 'builtins.range', 'range', (['runs'], {}), '(runs)\n', (6859, 6865), False, 'from builtins import range\n'), ((7201, 7222), 'numpy.random.seed', 'np.random.seed', (['(56467)'], {}), '(56467)\n', (7215, 7222), True, 'import numpy as np\n'), ((523, 561), 'scipy.stats.binom_test', 'ss.binom_test', (['ff', 'runs', 'expect_p_fail'], {}), '(ff, runs, expect_p_fail)\n', (536, 561), True, 'import scipy.stats as ss\n'), ((597, 658), 'scipy.stats.binom_test', 'ss.binom_test', (['ff', 'runs', 'expect_p_fail'], {'alternative': '"""greater"""'}), "(ff, runs, expect_p_fail, alternative='greater')\n", (610, 658), True, 'import scipy.stats as ss\n'), ((838, 857), 'numpy.min', 'np.min', (['pvals_2side'], {}), '(pvals_2side)\n', (844, 857), True, 'import numpy as np\n'), ((918, 937), 'numpy.min', 'np.min', (['pvals_1side'], {}), '(pvals_1side)\n', (924, 937), True, 'import numpy as np\n'), ((1332, 1344), 'mlpaper.util.area', 'area', (['*curve'], {}), '(*curve)\n', (1336, 1344), False, 'from mlpaper.util import area, interp1d\n'), ((1418, 1430), 'mlpaper.util.area', 'area', (['*curve'], {}), '(*curve)\n', (1422, 1430), False, 'from mlpaper.util import area, interp1d\n'), ((1454, 1471), 'mlpaper.util.area', 'area', (['*true_curve'], {}), '(*true_curve)\n', (1458, 1471), False, 'from mlpaper.util import area, interp1d\n'), ((1481, 1501), 'numpy.random.seed', 'np.random.seed', (['seed'], {}), '(seed)\n', (1495, 1501), True, 'import numpy as np\n'), ((1536, 1639), 'mlpaper.classification.curve_boot', 'curve_boot', (['y_true', 'y_score'], {'ref': 'true_value', 'curve_f': 'curve_f', 'confidence': 'confidence', 'x_grid': 'x_grid'}), '(y_true, y_score, ref=true_value, curve_f=curve_f, confidence=\n    confidence, x_grid=x_grid)\n', (1546, 1639), False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((1684, 1729), 'mlpaper.util.interp1d', 'interp1d', (['curve[cc.XGRID].values', '*true_curve'], {}), '(curve[cc.XGRID].values, *true_curve)\n', (1692, 1729), False, 'from mlpaper.util import area, interp1d\n'), ((2153, 2173), 'numpy.random.seed', 'np.random.seed', (['seed'], {}), '(seed)\n', (2167, 2173), True, 'import numpy as np\n'), ((2210, 2333), 'mlpaper.classification.curve_boot', 'curve_boot', (['y_true', 'y_score'], {'ref': 'y_score_ref', 'curve_f': 'curve_f', 'confidence': 'confidence', 'pairwise_CI': '(False)', 'x_grid': 'x_grid'}), '(y_true, y_score, ref=y_score_ref, curve_f=curve_f, confidence=\n    confidence, pairwise_CI=False, x_grid=x_grid)\n', (2220, 2333), False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((2537, 2574), 'numpy.all', 'np.all', (['(curve_.values == curve.values)'], {}), '(curve_.values == curve.values)\n', (2543, 2574), True, 'import numpy as np\n'), ((2584, 2604), 'numpy.random.seed', 'np.random.seed', (['seed'], {}), '(seed)\n', (2598, 2604), True, 'import numpy as np\n'), ((2641, 2763), 'mlpaper.classification.curve_boot', 'curve_boot', (['y_true', 'y_score'], {'ref': 'y_score_ref', 'curve_f': 'curve_f', 'confidence': 'confidence', 'pairwise_CI': '(True)', 'x_grid': 'x_grid'}), '(y_true, y_score, ref=y_score_ref, curve_f=curve_f, confidence=\n    confidence, pairwise_CI=True, x_grid=x_grid)\n', (2651, 2763), False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((2975, 3012), 'numpy.all', 'np.all', (['(curve_.values == curve.values)'], {}), '(curve_.values == curve.values)\n', (2981, 3012), True, 'import numpy as np\n'), ((3304, 3322), 'numpy.random.randn', 'np.random.randn', (['(2)'], {}), '(2)\n', (3319, 3322), True, 'import numpy as np\n'), ((3335, 3356), 'numpy.random.randn', 'np.random.randn', (['(2)', '(2)'], {}), '(2, 2)\n', (3350, 3356), True, 'import numpy as np\n'), ((3369, 3383), 'numpy.dot', 'np.dot', (['S', 'S.T'], {}), '(S, S.T)\n', (3375, 3383), True, 'import numpy as np\n'), ((3497, 3534), 'numpy.linspace', 'np.linspace', (['(0.0)', '(0.99)', 'DEFAULT_NGRID'], {}), '(0.0, 0.99, DEFAULT_NGRID)\n', (3508, 3534), True, 'import numpy as np\n'), ((3674, 3718), 'numpy.random.multivariate_normal', 'np.random.multivariate_normal', (['mu', 'S'], {'size': 'N'}), '(mu, S, size=N)\n', (3703, 3718), True, 'import numpy as np\n'), ((4142, 4180), 'numpy.linspace', 'np.linspace', (['(0.05)', '(0.95)', 'DEFAULT_NGRID'], {}), '(0.05, 0.95, DEFAULT_NGRID)\n', (4153, 4180), True, 'import numpy as np\n'), ((4909, 4946), 'numpy.linspace', 'np.linspace', (['(0.0)', '(0.99)', 'DEFAULT_NGRID'], {}), '(0.0, 0.99, DEFAULT_NGRID)\n', (4920, 4946), True, 'import numpy as np\n'), ((5842, 5859), 'numpy.random.randn', 'np.random.randn', ([], {}), '()\n', (5857, 5859), True, 'import numpy as np\n'), ((5956, 5966), 'numpy.mean', 'np.mean', (['x'], {}), '(x)\n', (5963, 5966), True, 'import numpy as np\n'), ((5980, 6010), 'mlpaper.mlpaper.boot_EB', 'bt.boot_EB', (['x'], {'confidence': '(0.95)'}), '(x, confidence=0.95)\n', (5990, 6010), True, 'import mlpaper.mlpaper as bt\n'), ((6331, 6348), 'numpy.random.randn', 'np.random.randn', ([], {}), '()\n', (6346, 6348), True, 'import numpy as np\n'), ((6439, 6501), 'mlpaper.mlpaper._boot_EB_and_test', 'bt._boot_EB_and_test', (['x'], {'confidence': 'confidence', 'return_CI': '(True)'}), '(x, confidence=confidence, return_CI=True)\n', (6459, 6501), True, 'import mlpaper.mlpaper as bt\n'), ((6602, 6677), 'mlpaper.mlpaper._boot_EB_and_test', 'bt._boot_EB_and_test', (['(x - true_value)'], {'confidence': 'confidence', 'return_CI': '(True)'}), '(x - true_value, confidence=confidence, return_CI=True)\n', (6622, 6677), True, 'import mlpaper.mlpaper as bt\n'), ((1775, 1799), 'numpy.abs', 'np.abs', (['(auc - true_value)'], {}), '(auc - true_value)\n', (1781, 1799), True, 'import numpy as np\n'), ((2103, 2143), 'numpy.all', 'np.all', (['(curve[cc.XGRID].values == x_grid)'], {}), '(curve[cc.XGRID].values == x_grid)\n', (2109, 2143), True, 'import numpy as np\n'), ((2827, 2848), 'numpy.abs', 'np.abs', (['(auc - auc_ref)'], {}), '(auc - auc_ref)\n', (2833, 2848), True, 'import numpy as np\n'), ((3557, 3579), 'numpy.array', 'np.array', (['[[0.0, 1.0]]'], {}), '([[0.0, 1.0]])\n', (3565, 3579), True, 'import numpy as np\n'), ((3581, 3603), 'numpy.array', 'np.array', (['[[0.0, 1.0]]'], {}), '([[0.0, 1.0]])\n', (3589, 3603), True, 'import numpy as np\n'), ((3633, 3650), 'numpy.random.rand', 'np.random.rand', (['N'], {}), '(N)\n', (3647, 3650), True, 'import numpy as np\n'), ((3730, 3747), 'numpy.random.randn', 'np.random.randn', ([], {}), '()\n', (3745, 3747), True, 'import numpy as np\n'), ((3800, 3841), 'numpy.random.choice', 'np.random.choice', (['N'], {'size': 'N', 'replace': '(True)'}), '(N, size=N, replace=True)\n', (3816, 3841), True, 'import numpy as np\n'), ((4515, 4537), 'numpy.array', 'np.array', (['[[0.0, 1.0]]'], {}), '([[0.0, 1.0]])\n', (4523, 4537), True, 'import numpy as np\n'), ((4539, 4557), 'numpy.array', 'np.array', (['[[p, p]]'], {}), '([[p, p]])\n', (4547, 4557), True, 'import numpy as np\n'), ((4969, 4991), 'numpy.array', 'np.array', (['[[0.0, 1.0]]'], {}), '([[0.0, 1.0]])\n', (4977, 4991), True, 'import numpy as np\n'), ((4993, 5015), 'numpy.array', 'np.array', (['[[0.0, 0.0]]'], {}), '([[0.0, 0.0]])\n', (5001, 5015), True, 'import numpy as np\n'), ((5879, 5896), 'numpy.random.randn', 'np.random.randn', ([], {}), '()\n', (5894, 5896), True, 'import numpy as np\n'), ((6028, 6047), 'numpy.abs', 'np.abs', (['(mu - mu_est)'], {}), '(mu - mu_est)\n', (6034, 6047), True, 'import numpy as np\n'), ((3949, 3960), 'numpy.zeros', 'np.zeros', (['N'], {}), '(N)\n', (3957, 3960), True, 'import numpy as np\n'), ((4012, 4023), 'numpy.zeros', 'np.zeros', (['N'], {}), '(N)\n', (4020, 4023), True, 'import numpy as np\n'), ((5919, 5937), 'numpy.random.randn', 'np.random.randn', (['N'], {}), '(N)\n', (5934, 5937), True, 'import numpy as np\n'), ((6892, 6910), 'numpy.random.randn', 'np.random.randn', (['N'], {}), '(N)\n', (6907, 6910), True, 'import numpy as np\n')]