import numpy as np
import pytest

from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1
from mr_uplift.mr_uplift import MRUplift, get_t_data
from mr_uplift.keras_model_functionality import prepare_data_optimized_loss
import sys
import pandas as pd

class TestMRUplift(object):

    def test_get_t_data(self):

        num_obs_1 = 10
        num_obs_2 = 3

        test_1 = get_t_data(0, num_obs_1)
        test_2 = get_t_data(np.array([0, 1]), num_obs_2)

        test_1_values = np.zeros(num_obs_1).reshape(-1, 1)
        test_2_values = np.concatenate([np.zeros(num_obs_2).reshape(-1, 1),
                                        np.ones(num_obs_2).reshape(-1, 1)], axis=1)

        assert np.mean(test_1 == test_1_values) == 1
        assert np.mean(test_2 == test_2_values) == 1

    def test_model_mean_outputs(self):
        true_ATE = np.array([[0, 0], [1, .5]])
        rmse_tolerance = .05
        num_obs = 10000

        y_no_noise, x_no_noise, tmt_no_noise = get_no_noise_data(num_obs)

        uplift_model = MRUplift()
        uplift_model.fit(x_no_noise, y_no_noise, tmt_no_noise.reshape(-1, 1),
                         n_jobs=1)
        oos_ice = uplift_model.predict_ice(response_transformer = True)

        assert np.sqrt(np.mean((oos_ice.mean(axis=1) -true_ATE)**2)) < rmse_tolerance

    def test_model_pred_oos_shapes(self):
        num_obs = 1000

        y, x, t = get_simple_uplift_data(num_obs)

        t = np.concatenate([t.reshape(-1, 1),
        np.random.binomial(1, .5, num_obs).reshape(-1, 1)], axis=1)

        param_grid = dict(num_nodes=[8], dropout=[.1], activation=[
                'relu'], num_layers=[1], epochs=[1], batch_size=[1000])

        uplift_model = MRUplift()
        uplift_model.fit(x, y, t, param_grid = param_grid, n_jobs=1)

        x_1 = x.copy()
        x_1 = pd.DataFrame(x_1)
        x_1.columns = ['var_'+str(x) for x in range(x.shape[1])]

        y_1 = y.copy()
        y_1 = pd.DataFrame(y_1)
        y_1.columns = ['var_'+str(x) for x in range(y.shape[1])]

        uplift_model_named = MRUplift()
        uplift_model_named.fit(x_1, y_1, t, param_grid = param_grid, n_jobs=1)

        assert uplift_model.predict_ice().shape == (
            np.unique(t, axis=0).shape[0], num_obs * .7, y.shape[1])

        assert uplift_model.predict_ice(x=x).shape == (np.unique(t,axis=0).shape[0],
            num_obs,
            y.shape[1])

        assert uplift_model.get_erupt_curves()

        assert uplift_model.get_erupt_curves(x = x, y = y, t = t)

        assert uplift_model_named.get_erupt_curves()

    def test_model_pred_oos_shapes_single_col_tmt(self):
        num_obs = 1000

        y, x, t = get_simple_uplift_data(num_obs)

        t = t.reshape(-1, 1)

        param_grid = dict(num_nodes=[8], dropout=[.1], activation=[
                                  'relu'], num_layers=[1], epochs=[1], batch_size=[1000])

        uplift_model = MRUplift()
        uplift_model.fit(x, y, t, param_grid = param_grid, n_jobs=1)

        assert uplift_model.predict_ice().shape == (
            np.unique(t, axis=0).shape[0], num_obs * .7, y.shape[1])

        assert uplift_model.predict_ice(x=x).shape == (np.unique(t,axis=0).shape[0],
            num_obs,
            y.shape[1])

        assert uplift_model.get_erupt_curves()

        assert uplift_model.get_erupt_curves(x = x, y = y, t = t)

    def test_model_pred_oos_shapes_single_col_tmt_propensity(self):
        num_obs = 1000

        y, x, t = get_simple_uplift_data(num_obs)

        t = t.reshape(-1, 1)

        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[2],
                                epochs=[1], batch_size=[100],
                                  alpha = [.5], copy_several_times = [1])

        uplift_model = MRUplift()
        uplift_model.fit(x, y, t, param_grid = param_grid, n_jobs=1,
            optimized_loss = True, use_propensity = True)

        assert uplift_model.predict_ice().shape == (
            np.unique(t, axis=0).shape[0], num_obs * .7, y.shape[1])

        assert uplift_model.predict_ice(x=x).shape == (np.unique(t,axis=0).shape[0],
            num_obs,
            y.shape[1])

        assert uplift_model.get_erupt_curves()

        assert uplift_model.get_erupt_curves(x = x, y = y, t = t)


    def test_prepare_data_optimized_loss_one_col_tmt(self):
        num_obs = 1000

        y, x, t = get_simple_uplift_data(num_obs)
        t = t.reshape(len(t),1)

        unique_treatments = np.unique(t, axis = 0)

        masks = np.ones(num_obs).reshape(num_obs,1)

        x, utility_weights, missing_utility, missing_y_mat, masks, weights = prepare_data_optimized_loss(x,y,t, masks ,unique_treatments)

        assert(utility_weights.shape == (num_obs, y.shape[1]))
        assert(missing_y_mat.shape == (num_obs, unique_treatments.shape[0], y.shape[1]))
        for q in range(unique_treatments.shape[0]):
            assert( ((missing_utility[:,q]==0) == (missing_y_mat[:,q,0] == -999)).mean() ==1 )

    def test_prepare_data_optimized_loss_two_col_tmt(self):
        num_obs = 1000

        y, x, t = get_simple_uplift_data(num_obs)

        t = np.concatenate([t.reshape(-1, 1),
        np.random.binomial(1, .5, num_obs).reshape(-1, 1)], axis=1)
        unique_treatments = np.unique(t, axis = 0)
        masks = np.ones(num_obs*len(unique_treatments)).reshape(num_obs,len(unique_treatments))

        x, utility_weights, missing_utility, missing_y_mat, masks, weights = prepare_data_optimized_loss(x,y,t,masks, unique_treatments)

        assert(utility_weights.shape == (num_obs, y.shape[1]))
        assert(missing_y_mat.shape == (num_obs, unique_treatments.shape[0], y.shape[1]))
        for q in range(unique_treatments.shape[0]):
            assert( ((missing_utility[:,q]==0) == (missing_y_mat[:,q,0] == -999)).mean() ==1 )

    def test_model_optim_mean_outputs(self):
        true_ATE = np.array([[0, 0], [1, .5]])
        rmse_tolerance = .05
        num_obs = 10000
        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[2], epochs=[30], batch_size=[100])

        y_no_noise, x_no_noise, tmt_no_noise = get_no_noise_data(num_obs)

        uplift_model = MRUplift()
        uplift_model.fit(x_no_noise, y_no_noise, tmt_no_noise.reshape(-1, 1),
                         n_jobs=1, param_grid = param_grid, optimized_loss = False)
        oos_ice = uplift_model.predict_ice(response_transformer = True)

        assert np.sqrt(np.mean((oos_ice.mean(axis=1) - true_ATE)**2)) < rmse_tolerance

    def test_model_get_random_erupts(self):
        true_ATE = np.array([[0, 0], [1, .5]])
        rmse_tolerance = .05
        num_obs = 10000
        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[2], epochs=[30], batch_size=[100],
                                  alpha = [.5], copy_several_times = [2])

        y_no_noise, x_no_noise, tmt_no_noise = get_no_noise_data(num_obs)

        uplift_model = MRUplift()
        uplift_model.fit(x_no_noise, y_no_noise, tmt_no_noise.reshape(-1, 1),
                         n_jobs=1, param_grid = param_grid, optimized_loss = True)
        oos_re = uplift_model.get_random_erupts()

        uplift_model_propensity = MRUplift()
        uplift_model_propensity.fit(x_no_noise, y_no_noise, tmt_no_noise.reshape(-1, 1),
                         n_jobs=1, param_grid = param_grid,
                         optimized_loss = True, use_propensity = True)
        oos_re_propensity = uplift_model_propensity.get_random_erupts()


        assert oos_re['mean'].iloc[0] > 0
        assert oos_re_propensity['mean'].iloc[0] > 0

    def test_varimp(self):
        num_obs = 10000
        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[2], epochs=[30], batch_size=[100])

        y, x, t = get_simple_uplift_data(num_obs)

        uplift_model = MRUplift()
        uplift_model.fit(x, y, t.reshape(-1, 1),
                         n_jobs=1, param_grid = param_grid)
        varimp = uplift_model.permutation_varimp(objective_weights = np.array([.7,-.3,0]).reshape(1,-1))

        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[2], epochs=[30], batch_size=[100],
                                  alpha = [.5], copy_several_times = [2])

        uplift_model_propensity = MRUplift()
        uplift_model_propensity.fit(x, y, t.reshape(-1, 1),
                         n_jobs=1, param_grid = param_grid,
                         optimized_loss = True, use_propensity = True)
        varimp_propensity = uplift_model_propensity.permutation_varimp(objective_weights = np.array([.7,-.3,0]).reshape(1,-1))

        assert varimp['permutation_varimp_metric'].iloc[0]>varimp['permutation_varimp_metric'].iloc[1]
        assert varimp_propensity['permutation_varimp_metric'].iloc[0]>varimp_propensity['permutation_varimp_metric'].iloc[1]


    def test_model_propensity(self):
        num_obs = 10000
        TOLERANCE = .98
        y, x, t, rule_assignment = get_observational_uplift_data_1(num_obs)

        param_grid = dict(num_nodes=[8], dropout=[.1], activation=['relu'], num_layers=[1],
                                epochs=[20], batch_size=[512],
                                  alpha = [.9999,.99], copy_several_times = [1])

        uplift_model = MRUplift()
        uplift_model.fit(x, y[:,0].reshape(-1,1), t, param_grid = param_grid, n_jobs=1,
            optimized_loss = True, use_propensity = True, test_size = 0)
        uplift_model.best_params_net
        y_test, x_test, t_test, rule_assignment_test = get_observational_uplift_data_1(num_obs)


        experiment_groups = np.zeros(num_obs)+2
        experiment_groups[np.where(x_test[:,-2]<.5)[0]] = 1
        experiment_groups[np.where(x_test[:,-2]<.33)[0]] = 0
        experiment_groups[np.where(x_test[:,-1]>.8)[0]] = 3


        optim_treatments_no_cuttoff = uplift_model.predict_optimal_treatments(x = x_test, use_propensity_score_cutoff = False)
        optim_treatments_cuttoff = uplift_model.predict_optimal_treatments(x = x_test, use_propensity_score_cutoff = True)

        optim_treatments_cuttoff_cat = optim_treatments_cuttoff.argmax(axis = 1)
        optim_treatments_no_cuttoff_cat = optim_treatments_no_cuttoff.argmax(axis = 1)

        correct_tmts_1 = np.array([x in [0,1] for x in optim_treatments_cuttoff_cat[np.where(experiment_groups == 0)[0]] ]).mean()
        correct_tmts_2 = np.array([x in [1,2] for x in optim_treatments_cuttoff_cat[np.where(experiment_groups == 1)[0]] ]).mean()
        correct_tmts_3 = np.array([x in [0,2] for x in optim_treatments_cuttoff_cat[np.where(experiment_groups == 2)[0]] ]).mean()
        correct_tmts_4 = np.array([x in [0] for x in optim_treatments_cuttoff_cat[np.where(experiment_groups == 3)[0]] ]).mean()

        correct_tmts_experiment_groups_1 = ((optim_treatments_cuttoff_cat[np.where(experiment_groups == 1)[0]] == 1)  == x_test[np.where(experiment_groups == 1)[0],0]).mean()
        correct_tmts_no_cutoff =  np.mean((optim_treatments_no_cuttoff_cat==1    ) == x_test[:,0])

        assert correct_tmts_1>TOLERANCE
        assert correct_tmts_2>TOLERANCE
        assert correct_tmts_3>TOLERANCE
        assert correct_tmts_4>TOLERANCE

        assert correct_tmts_experiment_groups_1>TOLERANCE

        assert np.array_equal(optim_treatments_cuttoff_cat,optim_treatments_no_cuttoff_cat) is False
        assert correct_tmts_no_cutoff>TOLERANCE

[((17, 17, 17, 41), 'mr_uplift.mr_uplift.get_t_data', 'get_t_data', ({(17, 28, 17, 29): '0', (17, 31, 17, 40): 'num_obs_1'}, {}), '(0, num_obs_1)', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((28, 19, 28, 46), 'numpy.array', 'np.array', ({(28, 28, 28, 45): '[[0, 0], [1, 0.5]]'}, {}), '([[0, 0], [1, 0.5]])', True, 'import numpy as np\n'), ((32, 47, 32, 73), 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', ({(32, 65, 32, 72): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((34, 23, 34, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((44, 18, 44, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(44, 41, 44, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((52, 23, 52, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((56, 14, 56, 31), 'pandas.DataFrame', 'pd.DataFrame', ({(56, 27, 56, 30): 'x_1'}, {}), '(x_1)', True, 'import pandas as pd\n'), ((60, 14, 60, 31), 'pandas.DataFrame', 'pd.DataFrame', ({(60, 27, 60, 30): 'y_1'}, {}), '(y_1)', True, 'import pandas as pd\n'), ((63, 29, 63, 39), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((82, 18, 82, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(82, 41, 82, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((89, 23, 89, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((106, 18, 106, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(106, 41, 106, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((114, 23, 114, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((133, 18, 133, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(133, 41, 133, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((136, 28, 136, 50), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((140, 77, 140, 137), 'mr_uplift.keras_model_functionality.prepare_data_optimized_loss', 'prepare_data_optimized_loss', ({(140, 105, 140, 106): 'x', (140, 107, 140, 108): 'y', (140, 109, 140, 110): 't', (140, 112, 140, 117): 'masks', (140, 119, 140, 136): 'unique_treatments'}, {}), '(x, y, t, masks, unique_treatments)', False, 'from mr_uplift.keras_model_functionality import prepare_data_optimized_loss\n'), ((150, 18, 150, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(150, 41, 150, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((154, 28, 154, 50), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((157, 77, 157, 136), 'mr_uplift.keras_model_functionality.prepare_data_optimized_loss', 'prepare_data_optimized_loss', ({(157, 105, 157, 106): 'x', (157, 107, 157, 108): 'y', (157, 109, 157, 110): 't', (157, 111, 157, 116): 'masks', (157, 118, 157, 135): 'unique_treatments'}, {}), '(x, y, t, masks, unique_treatments)', False, 'from mr_uplift.keras_model_functionality import prepare_data_optimized_loss\n'), ((165, 19, 165, 46), 'numpy.array', 'np.array', ({(165, 28, 165, 45): '[[0, 0], [1, 0.5]]'}, {}), '([[0, 0], [1, 0.5]])', True, 'import numpy as np\n'), ((170, 47, 170, 73), 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', ({(170, 65, 170, 72): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((172, 23, 172, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((180, 19, 180, 46), 'numpy.array', 'np.array', ({(180, 28, 180, 45): '[[0, 0], [1, 0.5]]'}, {}), '([[0, 0], [1, 0.5]])', True, 'import numpy as np\n'), ((186, 47, 186, 73), 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', ({(186, 65, 186, 72): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((188, 23, 188, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((193, 34, 193, 44), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((207, 18, 207, 49), 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', ({(207, 41, 207, 48): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((209, 23, 209, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((217, 34, 217, 44), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((230, 35, 230, 75), 'mr_uplift.dataset.data_simulation.get_observational_uplift_data_1', 'get_observational_uplift_data_1', ({(230, 67, 230, 74): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((236, 23, 236, 33), 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ({}, {}), '()', False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), ((240, 55, 240, 95), 'mr_uplift.dataset.data_simulation.get_observational_uplift_data_1', 'get_observational_uplift_data_1', ({(240, 87, 240, 94): 'num_obs'}, {}), '(num_obs)', False, 'from mr_uplift.dataset.data_simulation import get_no_noise_data, get_simple_uplift_data, get_observational_uplift_data_1\n'), ((261, 34, 261, 98), 'numpy.mean', 'np.mean', ({(261, 42, 261, 97): '(optim_treatments_no_cuttoff_cat == 1) == x_test[:, (0)]'}, {}), '((optim_treatments_no_cuttoff_cat == 1) == x_test[:, (0)])', True, 'import numpy as np\n'), ((18, 28, 18, 44), 'numpy.array', 'np.array', ({(18, 37, 18, 43): '[0, 1]'}, {}), '([0, 1])', True, 'import numpy as np\n'), ((24, 15, 24, 47), 'numpy.mean', 'np.mean', ({(24, 23, 24, 46): '(test_1 == test_1_values)'}, {}), '(test_1 == test_1_values)', True, 'import numpy as np\n'), ((25, 15, 25, 47), 'numpy.mean', 'np.mean', ({(25, 23, 25, 46): '(test_2 == test_2_values)'}, {}), '(test_2 == test_2_values)', True, 'import numpy as np\n'), ((243, 28, 243, 45), 'numpy.zeros', 'np.zeros', ({(243, 37, 243, 44): 'num_obs'}, {}), '(num_obs)', True, 'import numpy as np\n'), ((270, 15, 270, 91), 'numpy.array_equal', 'np.array_equal', ({(270, 30, 270, 58): 'optim_treatments_cuttoff_cat', (270, 59, 270, 90): 'optim_treatments_no_cuttoff_cat'}, {}), '(optim_treatments_cuttoff_cat, optim_treatments_no_cuttoff_cat)', True, 'import numpy as np\n'), ((20, 24, 20, 43), 'numpy.zeros', 'np.zeros', ({(20, 33, 20, 42): 'num_obs_1'}, {}), '(num_obs_1)', True, 'import numpy as np\n'), ((138, 16, 138, 32), 'numpy.ones', 'np.ones', ({(138, 24, 138, 31): 'num_obs'}, {}), '(num_obs)', True, 'import numpy as np\n'), ((244, 26, 244, 51), 'numpy.where', 'np.where', ({(244, 35, 244, 50): '(x_test[:, (-2)] < 0.5)'}, {}), '(x_test[:, (-2)] < 0.5)', True, 'import numpy as np\n'), ((245, 26, 245, 52), 'numpy.where', 'np.where', ({(245, 35, 245, 51): '(x_test[:, (-2)] < 0.33)'}, {}), '(x_test[:, (-2)] < 0.33)', True, 'import numpy as np\n'), ((246, 26, 246, 51), 'numpy.where', 'np.where', ({(246, 35, 246, 50): '(x_test[:, (-1)] > 0.8)'}, {}), '(x_test[:, (-1)] > 0.8)', True, 'import numpy as np\n'), ((21, 40, 21, 59), 'numpy.zeros', 'np.zeros', ({(21, 49, 21, 58): 'num_obs_2'}, {}), '(num_obs_2)', True, 'import numpy as np\n'), ((22, 40, 22, 58), 'numpy.ones', 'np.ones', ({(22, 48, 22, 57): 'num_obs_2'}, {}), '(num_obs_2)', True, 'import numpy as np\n'), ((47, 8, 47, 42), 'numpy.random.binomial', 'np.random.binomial', ({(47, 27, 47, 28): '1', (47, 30, 47, 32): '0.5', (47, 34, 47, 41): 'num_obs'}, {}), '(1, 0.5, num_obs)', True, 'import numpy as np\n'), ((67, 12, 67, 32), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((69, 55, 69, 74), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((93, 12, 93, 32), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((95, 55, 95, 74), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((119, 12, 119, 32), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((121, 55, 121, 74), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((153, 8, 153, 42), 'numpy.random.binomial', 'np.random.binomial', ({(153, 27, 153, 28): '1', (153, 30, 153, 32): '0.5', (153, 34, 153, 41): 'num_obs'}, {}), '(1, 0.5, num_obs)', True, 'import numpy as np\n'), ((212, 69, 212, 89), 'numpy.array', 'np.array', ({(212, 78, 212, 88): '[0.7, -0.3, 0]'}, {}), '([0.7, -0.3, 0])', True, 'import numpy as np\n'), ((221, 91, 221, 111), 'numpy.array', 'np.array', ({(221, 100, 221, 110): '[0.7, -0.3, 0]'}, {}), '([0.7, -0.3, 0])', True, 'import numpy as np\n'), ((260, 74, 260, 106), 'numpy.where', 'np.where', ({(260, 83, 260, 105): 'experiment_groups == 1'}, {}), '(experiment_groups == 1)', True, 'import numpy as np\n'), ((260, 128, 260, 160), 'numpy.where', 'np.where', ({(260, 137, 260, 159): 'experiment_groups == 1'}, {}), '(experiment_groups == 1)', True, 'import numpy as np\n'), ((255, 84, 255, 116), 'numpy.where', 'np.where', ({(255, 93, 255, 115): 'experiment_groups == 0'}, {}), '(experiment_groups == 0)', True, 'import numpy as np\n'), ((256, 84, 256, 116), 'numpy.where', 'np.where', ({(256, 93, 256, 115): 'experiment_groups == 1'}, {}), '(experiment_groups == 1)', True, 'import numpy as np\n'), ((257, 84, 257, 116), 'numpy.where', 'np.where', ({(257, 93, 257, 115): 'experiment_groups == 2'}, {}), '(experiment_groups == 2)', True, 'import numpy as np\n'), ((258, 82, 258, 114), 'numpy.where', 'np.where', ({(258, 91, 258, 113): 'experiment_groups == 3'}, {}), '(experiment_groups == 3)', True, 'import numpy as np\n')]