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, 'mr_uplift.mr_uplift.get_t_data', 'get_t_data', (['(0)', 'num_obs_1'], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (28, 'numpy.array', 'np.array', (['[[0, 0], [1, 0.5]]'], {}), True, 'import numpy as np\n'), (32, 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (44, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (56, 'pandas.DataFrame', 'pd.DataFrame', (['x_1'], {}), True, 'import pandas as pd\n'), (60, 'pandas.DataFrame', 'pd.DataFrame', (['y_1'], {}), True, 'import pandas as pd\n'), (63, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (82, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (106, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (133, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (140, 'mr_uplift.keras_model_functionality.prepare_data_optimized_loss', 'prepare_data_optimized_loss', (['x', 'y', 't', 'masks', 'unique_treatments'], {}), False, 'from mr_uplift.keras_model_functionality import prepare_data_optimized_loss\n'), (150, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (157, 'mr_uplift.keras_model_functionality.prepare_data_optimized_loss', 'prepare_data_optimized_loss', (['x', 'y', 't', 'masks', 'unique_treatments'], {}), False, 'from mr_uplift.keras_model_functionality import prepare_data_optimized_loss\n'), (165, 'numpy.array', 'np.array', (['[[0, 0], [1, 0.5]]'], {}), True, 'import numpy as np\n'), (170, 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (180, 'numpy.array', 'np.array', (['[[0, 0], [1, 0.5]]'], {}), True, 'import numpy as np\n'), (186, 'mr_uplift.dataset.data_simulation.get_no_noise_data', 'get_no_noise_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (193, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (207, 'mr_uplift.dataset.data_simulation.get_simple_uplift_data', 'get_simple_uplift_data', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (217, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (230, 'mr_uplift.dataset.data_simulation.get_observational_uplift_data_1', 'get_observational_uplift_data_1', (['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, 'mr_uplift.mr_uplift.MRUplift', 'MRUplift', ([], {}), False, 'from mr_uplift.mr_uplift import MRUplift, get_t_data\n'), (240, 'mr_uplift.dataset.data_simulation.get_observational_uplift_data_1', 'get_observational_uplift_data_1', (['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, 'numpy.mean', 'np.mean', (['((optim_treatments_no_cuttoff_cat == 1) == x_test[:, (0)])'], {}), True, 'import numpy as np\n'), (18, 'numpy.array', 'np.array', (['[0, 1]'], {}), True, 'import numpy as np\n'), (24, 'numpy.mean', 'np.mean', (['(test_1 == test_1_values)'], {}), True, 'import numpy as np\n'), (25, 'numpy.mean', 'np.mean', (['(test_2 == test_2_values)'], {}), True, 'import numpy as np\n'), (243, 'numpy.zeros', 'np.zeros', (['num_obs'], {}), True, 'import numpy as np\n'), (270, 'numpy.array_equal', 'np.array_equal', (['optim_treatments_cuttoff_cat', 'optim_treatments_no_cuttoff_cat'], {}), True, 'import numpy as np\n'), (20, 'numpy.zeros', 'np.zeros', (['num_obs_1'], {}), True, 'import numpy as np\n'), (138, 'numpy.ones', 'np.ones', (['num_obs'], {}), True, 'import numpy as np\n'), (244, 'numpy.where', 'np.where', (['(x_test[:, (-2)] < 0.5)'], {}), True, 'import numpy as np\n'), (245, 'numpy.where', 'np.where', (['(x_test[:, (-2)] < 0.33)'], {}), True, 'import numpy as np\n'), (246, 'numpy.where', 'np.where', (['(x_test[:, (-1)] > 0.8)'], {}), True, 'import numpy as np\n'), (21, 'numpy.zeros', 'np.zeros', (['num_obs_2'], {}), True, 'import numpy as np\n'), (22, 'numpy.ones', 'np.ones', (['num_obs_2'], {}), True, 'import numpy as np\n'), (47, 'numpy.random.binomial', 'np.random.binomial', (['(1)', '(0.5)', 'num_obs'], {}), True, 'import numpy as np\n'), (67, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (69, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (93, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (95, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (119, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (121, 'numpy.unique', 'np.unique', (['t'], {'axis': '(0)'}), True, 'import numpy as np\n'), (153, 'numpy.random.binomial', 'np.random.binomial', (['(1)', '(0.5)', 'num_obs'], {}), True, 'import numpy as np\n'), (212, 'numpy.array', 'np.array', (['[0.7, -0.3, 0]'], {}), True, 'import numpy as np\n'), (221, 'numpy.array', 'np.array', (['[0.7, -0.3, 0]'], {}), True, 'import numpy as np\n'), (260, 'numpy.where', 'np.where', (['(experiment_groups == 1)'], {}), True, 'import numpy as np\n'), (260, 'numpy.where', 'np.where', (['(experiment_groups == 1)'], {}), True, 'import numpy as np\n'), (255, 'numpy.where', 'np.where', (['(experiment_groups == 0)'], {}), True, 'import numpy as np\n'), (256, 'numpy.where', 'np.where', (['(experiment_groups == 1)'], {}), True, 'import numpy as np\n'), (257, 'numpy.where', 'np.where', (['(experiment_groups == 2)'], {}), True, 'import numpy as np\n'), (258, 'numpy.where', 'np.where', (['(experiment_groups == 3)'], {}), True, 'import numpy as np\n')]