# 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")

[((85, 21, 85, 55), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((86, 20, 86, 54), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((87, 21, 87, 55), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((88, 14, 88, 25), 'builtins.range', 'range', ({(88, 20, 88, 24): 'runs'}, {}), '(runs)', False, 'from builtins import range\n'), ((153, 14, 153, 25), 'builtins.range', 'range', ({(153, 20, 153, 24): 'runs'}, {}), '(runs)', False, 'from builtins import range\n'), ((169, 9, 169, 26), 'numpy.random.randn', 'np.random.randn', ({}, {}), '()', True, 'import numpy as np\n'), ((187, 14, 187, 25), 'builtins.range', 'range', ({(187, 20, 187, 24): 'runs'}, {}), '(runs)', False, 'from builtins import range\n'), ((200, 4, 200, 25), 'numpy.random.seed', 'np.random.seed', ({(200, 19, 200, 24): '(56467)'}, {}), '(56467)', True, 'import numpy as np\n'), ((20, 19, 20, 57), 'scipy.stats.binom_test', 'ss.binom_test', ({(20, 33, 20, 35): 'ff', (20, 37, 20, 41): 'runs', (20, 43, 20, 56): 'expect_p_fail'}, {}), '(ff, runs, expect_p_fail)', True, 'import scipy.stats as ss\n'), ((21, 19, 21, 80), 'scipy.stats.binom_test', 'ss.binom_test', (), '', True, 'import scipy.stats as ss\n'), ((26, 11, 26, 30), 'numpy.min', 'np.min', ({(26, 18, 26, 29): 'pvals_2side'}, {}), '(pvals_2side)', True, 'import numpy as np\n'), ((28, 11, 28, 30), 'numpy.min', 'np.min', ({(28, 18, 28, 29): 'pvals_1side'}, {}), '(pvals_1side)', True, 'import numpy as np\n'), ((42, 15, 42, 27), 'mlpaper.util.area', 'area', ({(42, 20, 42, 26): '*curve'}, {}), '(*curve)', False, 'from mlpaper.util import area, interp1d\n'), ((44, 19, 44, 31), 'mlpaper.util.area', 'area', ({(44, 24, 44, 30): '*curve'}, {}), '(*curve)', False, 'from mlpaper.util import area, interp1d\n'), ((46, 22, 46, 39), 'mlpaper.util.area', 'area', ({(46, 27, 46, 38): '*true_curve'}, {}), '(*true_curve)', False, 'from mlpaper.util import area, interp1d\n'), ((48, 8, 48, 28), 'numpy.random.seed', 'np.random.seed', ({(48, 23, 48, 27): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((49, 34, 51, 9), 'mlpaper.classification.curve_boot', 'curve_boot', (), '', False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((52, 27, 52, 72), 'mlpaper.util.interp1d', 'interp1d', ({(52, 36, 52, 58): 'curve[cc.XGRID].values', (52, 60, 52, 71): '*true_curve'}, {}), '(curve[cc.XGRID].values, *true_curve)', False, 'from mlpaper.util import area, interp1d\n'), ((62, 8, 62, 28), 'numpy.random.seed', 'np.random.seed', ({(62, 23, 62, 27): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((63, 36, 65, 9), 'mlpaper.classification.curve_boot', 'curve_boot', (), '', False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((70, 15, 70, 52), 'numpy.all', 'np.all', ({(70, 22, 70, 51): '(curve_.values == curve.values)'}, {}), '(curve_.values == curve.values)', True, 'import numpy as np\n'), ((72, 8, 72, 28), 'numpy.random.seed', 'np.random.seed', ({(72, 23, 72, 27): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((73, 36, 75, 9), 'mlpaper.classification.curve_boot', 'curve_boot', (), '', False, 'from mlpaper.classification import DEFAULT_NGRID, curve_boot\n'), ((80, 15, 80, 52), 'numpy.all', 'np.all', ({(80, 22, 80, 51): '(curve_.values == curve.values)'}, {}), '(curve_.values == curve.values)', True, 'import numpy as np\n'), ((89, 13, 89, 31), 'numpy.random.randn', 'np.random.randn', ({(89, 29, 89, 30): '2'}, {}), '(2)', True, 'import numpy as np\n'), ((90, 12, 90, 33), 'numpy.random.randn', 'np.random.randn', ({(90, 28, 90, 29): '2', (90, 31, 90, 32): '2'}, {}), '(2, 2)', True, 'import numpy as np\n'), ((91, 12, 91, 26), 'numpy.dot', 'np.dot', ({(91, 19, 91, 20): 'S', (91, 22, 91, 25): 'S.T'}, {}), '(S, S.T)', True, 'import numpy as np\n'), ((95, 17, 95, 54), 'numpy.linspace', 'np.linspace', ({(95, 29, 95, 32): '0.0', (95, 34, 95, 38): '0.99', (95, 40, 95, 53): 'DEFAULT_NGRID'}, {}), '(0.0, 0.99, DEFAULT_NGRID)', True, 'import numpy as np\n'), ((98, 18, 98, 62), 'numpy.random.multivariate_normal', 'np.random.multivariate_normal', (), '', True, 'import numpy as np\n'), ((107, 17, 107, 55), 'numpy.linspace', 'np.linspace', ({(107, 29, 107, 33): '0.05', (107, 35, 107, 39): '0.95', (107, 41, 107, 54): 'DEFAULT_NGRID'}, {}), '(0.05, 0.95, DEFAULT_NGRID)', True, 'import numpy as np\n'), ((127, 17, 127, 54), 'numpy.linspace', 'np.linspace', ({(127, 29, 127, 32): '0.0', (127, 34, 127, 38): '0.99', (127, 40, 127, 53): 'DEFAULT_NGRID'}, {}), '(0.0, 0.99, DEFAULT_NGRID)', True, 'import numpy as np\n'), ((154, 13, 154, 30), 'numpy.random.randn', 'np.random.randn', ({}, {}), '()', True, 'import numpy as np\n'), ((158, 17, 158, 27), 'numpy.mean', 'np.mean', ({(158, 25, 158, 26): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((159, 13, 159, 43), 'mlpaper.mlpaper.boot_EB', 'bt.boot_EB', (), '', True, 'import mlpaper.mlpaper as bt\n'), ((170, 19, 170, 36), 'numpy.random.randn', 'np.random.randn', ({}, {}), '()', True, 'import numpy as np\n'), ((176, 19, 176, 81), 'mlpaper.mlpaper._boot_EB_and_test', 'bt._boot_EB_and_test', (), '', True, 'import mlpaper.mlpaper as bt\n'), ((180, 22, 180, 97), 'mlpaper.mlpaper._boot_EB_and_test', 'bt._boot_EB_and_test', (), '', True, 'import mlpaper.mlpaper as bt\n'), ((54, 18, 54, 42), 'numpy.abs', 'np.abs', ({(54, 25, 54, 41): '(auc - true_value)'}, {}), '(auc - true_value)', True, 'import numpy as np\n'), ((60, 35, 60, 75), 'numpy.all', 'np.all', ({(60, 42, 60, 74): '(curve[cc.XGRID].values == x_grid)'}, {}), '(curve[cc.XGRID].values == x_grid)', True, 'import numpy as np\n'), ((77, 19, 77, 40), 'numpy.abs', 'np.abs', ({(77, 26, 77, 39): '(auc - auc_ref)'}, {}), '(auc - auc_ref)', True, 'import numpy as np\n'), ((96, 22, 96, 44), 'numpy.array', 'np.array', ({(96, 31, 96, 43): '[[0.0, 1.0]]'}, {}), '([[0.0, 1.0]])', True, 'import numpy as np\n'), ((96, 46, 96, 68), 'numpy.array', 'np.array', ({(96, 55, 96, 67): '[[0.0, 1.0]]'}, {}), '([[0.0, 1.0]])', True, 'import numpy as np\n'), ((97, 17, 97, 34), 'numpy.random.rand', 'np.random.rand', ({(97, 32, 97, 33): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((99, 11, 99, 28), 'numpy.random.randn', 'np.random.randn', ({}, {}), '()', True, 'import numpy as np\n'), ((100, 18, 100, 59), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((117, 22, 117, 44), 'numpy.array', 'np.array', ({(117, 31, 117, 43): '[[0.0, 1.0]]'}, {}), '([[0.0, 1.0]])', True, 'import numpy as np\n'), ((117, 46, 117, 64), 'numpy.array', 'np.array', ({(117, 55, 117, 63): '[[p, p]]'}, {}), '([[p, p]])', True, 'import numpy as np\n'), ((128, 22, 128, 44), 'numpy.array', 'np.array', ({(128, 31, 128, 43): '[[0.0, 1.0]]'}, {}), '([[0.0, 1.0]])', True, 'import numpy as np\n'), ((128, 46, 128, 68), 'numpy.array', 'np.array', ({(128, 55, 128, 67): '[[0.0, 0.0]]'}, {}), '([[0.0, 0.0]])', True, 'import numpy as np\n'), ((155, 19, 155, 36), 'numpy.random.randn', 'np.random.randn', ({}, {}), '()', True, 'import numpy as np\n'), ((161, 16, 161, 35), 'numpy.abs', 'np.abs', ({(161, 23, 161, 34): '(mu - mu_est)'}, {}), '(mu - mu_est)', True, 'import numpy as np\n'), ((103, 28, 103, 39), 'numpy.zeros', 'np.zeros', ({(103, 37, 103, 38): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((104, 32, 104, 43), 'numpy.zeros', 'np.zeros', ({(104, 41, 104, 42): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((156, 21, 156, 39), 'numpy.random.randn', 'np.random.randn', ({(156, 37, 156, 38): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((188, 25, 188, 43), 'numpy.random.randn', 'np.random.randn', ({(188, 41, 188, 42): 'N'}, {}), '(N)', True, 'import numpy as np\n')]