from __future__ import division # so that, under Python 2.*, integer division results in real numbers

import unittest
import warnings

import numpy as np

from expan.core.results import CombinedTestStatistics
from expan.core.statistical_test import *
from expan.core.experiment import Experiment
from expan.core.util import generate_random_data, find_value_by_key_with_condition


class ExperimentTestCase(unittest.TestCase):
    """ Defines the setUp() and tearDown() functions for the statistics test cases."""

    def setUp(self):
        """ Load the needed data sets for all StatisticsTestCases and set the random
        seed so that randomized algorithms show deterministic behaviour."""
        np.random.seed(0)
        data, metadata = generate_random_data()

        self.data, self.metadata = data, metadata

        feature_has = FeatureFilter('feature', 'has')
        feature_non = FeatureFilter('feature', 'non')
        feature_one = FeatureFilter('feature', 'feature that only has one data point')

        # simple statistical test
        self.kpi = KPI('normal_same')
        self.variants = Variants('variant', 'B', 'A')
        self.nonsense_variants = Variants('variant', 'C', 'D')
        self.test_normal_same = StatisticalTest(self.data, self.kpi, [], self.variants)
        self.test_nonsense_variant = StatisticalTest(self.data, self.kpi, [], self.nonsense_variants)
        self.test_normal_same_feature_non = StatisticalTest(self.data, self.kpi, [feature_non], self.variants)
        self.test_normal_same_feature_has = StatisticalTest(self.data, self.kpi, [feature_has], self.variants)
        self.test_normal_same_feature_one = StatisticalTest(self.data, self.kpi, [feature_one], self.variants)

        # statistical test with derived kpi
        self.derived_kpi = DerivedKPI('derived_kpi_one', 'normal_same', 'normal_shifted')
        self.test_derived_kpi = StatisticalTest(self.data, self.derived_kpi, [], self.variants)

        # small dummy data frame
        data_dummy = np.array([['index', 'entity', 'variant', 'normal_same', 'normal_shifted'],
                               [0, 1, 'A', 2.0, 1.0], [1, 2, 'B', 3.0, 2.0],
                               [2, 3, 'A', 3.0, 2.0], [3, 4, 'B', 2.5, 1.0]])
        self.data_dummy_df = pd.DataFrame(data=data_dummy[1:, 1:],
                                          columns=data_dummy[0, 1:]).convert_objects(convert_numeric=True)

        data_dummy_with_nan = np.array([['index', 'entity', 'variant', 'normal_same', 'normal_shifted'],
                               [0, 1, 'A', 2.0, 1.0], [1, 2, 'B', 3.0, 2.0],
                               [2, 3, 'A', 3.0, 2.0], [3, 4, 'B', 2.5, 1.0],
                               [4, 5, 'A', 0.0, 0.0], [5, 6, 'B', None, None]])
        self.data_dummy_df_with_nan = pd.DataFrame(data=data_dummy_with_nan[1:, 1:],
                                                   columns=data_dummy_with_nan[0, 1:]).convert_objects(convert_numeric=True)

        # statistical test suite
        self.suite_with_one_test = StatisticalTestSuite([self.test_normal_same])
        self.suite_with_one_test_correction = StatisticalTestSuite([self.test_normal_same], CorrectionMethod.BH)
        self.suite_with_two_tests = StatisticalTestSuite([self.test_normal_same, self.test_derived_kpi],
                                                         CorrectionMethod.BONFERRONI)
        self.suite_with_one_subgroup = StatisticalTestSuite([self.test_normal_same_feature_one])
        self.suite_with_three_subgroups = StatisticalTestSuite([self.test_normal_same_feature_non,
                                                                self.test_normal_same_feature_has,
                                                                self.test_normal_same_feature_one],
                                                               CorrectionMethod.BH)

        # small dummy data frames producing zero std and -1 statistical power
        data_dummy_zero_std = np.array([['index', 'entity', 'variant', 'normal_same', 'normal_shifted'],
                                        [0, 1, 'A', 1.0, 1.0], [1, 2, 'B', 2.0, 2.0],
                                        [2, 3, 'A', 1.0, 1.0], [3, 4, 'B', 2.0, 2.0]])
        self.data_dummy_zero_std = pd.DataFrame(data=data_dummy_zero_std[1:, 1:],
                                                columns=data_dummy_zero_std[0, 1:]).convert_objects(
            convert_numeric=True)
        self.test_normal_same_zero_std = StatisticalTest(self.data_dummy_zero_std, self.kpi, [], self.variants)
        self.suite_with_one_test_zero_std = StatisticalTestSuite([self.test_normal_same_zero_std])

        # small dummy data frames with all nan values
        self.data_dummy_all_nan = pd.DataFrame({
                        'entity' : [1,2,3,4],
                        'variant' : ['A','B','C','D'],
                        'normal_same' : [np.nan] * 4,
                        'normal_shifted' : [np.nan] * 4,
                            })
        self.test_normal_same_nan_data = StatisticalTest(self.data_dummy_all_nan, self.kpi, [], self.variants)
        self.suite_with_one_test_with_nan_data = StatisticalTestSuite([self.test_normal_same_nan_data])

    def tearDown(self):
        """ Clean up after the test """
        pass

    def getExperiment(self):
        return Experiment(self.metadata)


class ExperimentClassTestCases(ExperimentTestCase):

    def test_Experiment_constructor(self):
        experiment = self.getExperiment()
        self.assertTrue(isinstance(experiment.metadata, dict))


class StatisticalTestTestCases(ExperimentTestCase):
    """ Test the method analyze_statistical_test. """

    def test_fixed_horizon(self):
        ndecimals = 5
        res = self.getExperiment().analyze_statistical_test(self.test_normal_same, test_method='fixed_horizon')

        self.assertAlmostEqual(res.result.delta, 0.033053, ndecimals)

        lower_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 2.5, 'value')
        upper_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 97.5, 'value')
        self.assertAlmostEqual(lower_bound_ci, -0.007135, ndecimals)
        self.assertAlmostEqual(upper_bound_ci, 0.073240, ndecimals)

        self.assertEqual(res.result.treatment_statistics.sample_size, 6108)
        self.assertEqual(res.result.control_statistics.sample_size,   3892)

        self.assertAlmostEqual(res.result.treatment_statistics.mean,  0.025219, ndecimals)
        self.assertAlmostEqual(res.result.control_statistics.mean,   -0.007833, ndecimals)

        self.assertAlmostEqual(res.result.statistical_power, 0.36401, ndecimals)
        self.assertEqual(res.test.kpi.name, 'normal_same')

    def test_group_sequential(self):
        ndecimals = 5
        res = self.getExperiment().analyze_statistical_test(self.test_normal_same, test_method='group_sequential')

        self.assertAlmostEqual(res.result.delta, 0.033053, ndecimals)

        lower_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 2.5, 'value')
        upper_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 97.5, 'value')
        self.assertAlmostEqual(lower_bound_ci, -0.007135, ndecimals)
        self.assertAlmostEqual(upper_bound_ci, 0.073240, ndecimals)

        self.assertEqual(res.result.treatment_statistics.sample_size, 6108)
        self.assertEqual(res.result.control_statistics.sample_size, 3892)

        self.assertAlmostEqual(res.result.treatment_statistics.mean, 0.025219, ndecimals)
        self.assertAlmostEqual(res.result.control_statistics.mean, -0.007833, ndecimals)

        self.assertAlmostEqual(res.result.statistical_power, 0.36401, ndecimals)
        self.assertEqual(res.result.stop, False)
        self.assertEqual(res.test.kpi.name, 'normal_same')

    def test_bayes_factor_delta(self):
        ndecimals = 5
        res = self.getExperiment().analyze_statistical_test(self.test_normal_same,
                                                            test_method='bayes_factor',
                                                            num_iters=2000)

        self.assertAlmostEqual(res.result.delta, 0.033053, ndecimals)

        lower_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 2.5, 'value')
        upper_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 97.5, 'value')
        self.assertAlmostEqual(lower_bound_ci, -0.00829, ndecimals)
        self.assertAlmostEqual(upper_bound_ci, 0.07127, ndecimals)

        self.assertEqual(res.result.treatment_statistics.sample_size, 6108)
        self.assertEqual(res.result.control_statistics.sample_size, 3892)

        self.assertAlmostEqual(res.result.treatment_statistics.mean, 0.025219, ndecimals)
        self.assertAlmostEqual(res.result.control_statistics.mean, -0.007833, ndecimals)

        self.assertEqual(res.result.stop, True)
        self.assertEqual(res.test.kpi.name, 'normal_same')

    def test_bayes_precision_delta(self):
        ndecimals = 5
        res = self.getExperiment().analyze_statistical_test(self.test_normal_same,
                                                            test_method='bayes_precision',
                                                            num_iters=2000)
        self.assertAlmostEqual(res.result.delta, 0.033053, ndecimals)

        lower_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 2.5, 'value')
        upper_bound_ci = find_value_by_key_with_condition(res.result.confidence_interval, 'percentile', 97.5, 'value')
        self.assertAlmostEqual(lower_bound_ci, -0.00829, ndecimals)
        self.assertAlmostEqual(upper_bound_ci, 0.07127, ndecimals)

        self.assertEqual(res.result.treatment_statistics.sample_size, 6108)
        self.assertEqual(res.result.control_statistics.sample_size, 3892)

        self.assertAlmostEqual(res.result.treatment_statistics.mean, 0.025219, ndecimals)
        self.assertAlmostEqual(res.result.control_statistics.mean, -0.007833, ndecimals)

        self.assertEqual(res.result.stop, True)


class StatisticalTestSuiteTestCases(ExperimentTestCase):
    """ Test the method analyze_statistical_test_suite. """

    def test_one_test_in_suite(self):
        ndecimals = 5
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_one_test)

        self.assertEqual(res.correction_method, CorrectionMethod.NONE)
        self.assertEqual(len(res.results), 1)
        simple_stat_res = res.results[0].result

        self.assertAlmostEqual(simple_stat_res.original_test_statistics.delta, 0.033053, ndecimals)
        lower_bound_ci = find_value_by_key_with_condition(simple_stat_res.original_test_statistics.confidence_interval,
                                                          'percentile', 2.5, 'value')
        upper_bound_ci = find_value_by_key_with_condition(simple_stat_res.original_test_statistics.confidence_interval,
                                                          'percentile', 97.5, 'value')
        self.assertAlmostEqual(lower_bound_ci, -0.007135, ndecimals)
        self.assertAlmostEqual(upper_bound_ci, 0.073240, ndecimals)

        self.assertEqual(simple_stat_res.original_test_statistics.treatment_statistics.sample_size, 6108)
        self.assertEqual(simple_stat_res.original_test_statistics.control_statistics.sample_size,   3892)

        self.assertAlmostEqual(simple_stat_res.original_test_statistics.treatment_statistics.mean,  0.025219, ndecimals)
        self.assertAlmostEqual(simple_stat_res.original_test_statistics.control_statistics.mean,   -0.007833, ndecimals)
        self.assertAlmostEqual(simple_stat_res.original_test_statistics.statistical_power, 0.36401, ndecimals)

        # corrected test statistics should be the same as original test statistics (checks some values)
        self.assertAlmostEqual(simple_stat_res.corrected_test_statistics.delta, 0.033053, ndecimals)
        self.assertAlmostEqual(simple_stat_res.corrected_test_statistics.statistical_power, 0.36401, ndecimals)
        self.assertEqual(simple_stat_res.corrected_test_statistics.treatment_statistics.sample_size, 6108)
        self.assertEqual(simple_stat_res.corrected_test_statistics.control_statistics.sample_size, 3892)

    def test_one_test_in_suite_with_wrong_correction(self):
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_one_test_correction)
        self.assertEqual(res.correction_method, CorrectionMethod.NONE)
        self.assertEqual(len(res.results), 1)

    def test_two_tests_in_suite(self):
        ndecimals = 5

        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_two_tests)
        self.assertEqual(res.correction_method, CorrectionMethod.BONFERRONI)
        self.assertEqual(len(res.results), 2)

        res_normal_same = res.results[0]
        res_derived_kpi = res.results[1]

        self.assertEqual(res_normal_same.test.kpi.name, "normal_same")
        self.assertEqual(res_derived_kpi.test.kpi.name, "derived_kpi_one")
        self.assertTrue(isinstance(res_normal_same.result, CombinedTestStatistics))
        self.assertTrue(isinstance(res_derived_kpi.result, CombinedTestStatistics))

        power_normal_same_before = res_normal_same.result.original_test_statistics.statistical_power
        power_normal_same_corrected = res_normal_same.result.corrected_test_statistics.statistical_power
        self.assertAlmostEqual(power_normal_same_before, 0.36401, ndecimals)
        self.assertLess(power_normal_same_corrected, power_normal_same_before)

        power_derived_kpi_before = res_derived_kpi.result.original_test_statistics.statistical_power
        power_derived_kpi_corrected = res_derived_kpi.result.corrected_test_statistics.statistical_power
        self.assertAlmostEqual(power_derived_kpi_before, 0.3472, ndecimals)
        self.assertLess(power_derived_kpi_corrected, power_derived_kpi_before)

    def test_one_subgroup_in_suite(self):
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_one_subgroup)
        self.assertEqual(res.correction_method, CorrectionMethod.NONE)
        self.assertEqual(len(res.results), 0)

    def test_three_subgroup_in_suite(self):
        """ Two subgroups contain valid data. One subgroup contains only one entity.
        When there is no enough data in one subgroup, which means one subgroup might contain zero
        or only one variant. (e.g. "browser" = "hack name"), it should ignore this subgroup and continue
        doing analysis with other subgroups.
        """
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_three_subgroups)
        self.assertEqual(res.correction_method, CorrectionMethod.BH)
        self.assertEqual(len(res.results), 2)

        res_subgroup_feature_non = res.results[0]
        res_subgroup_feature_has = res.results[1]

        self.assertTrue(isinstance(res_subgroup_feature_non.result, CombinedTestStatistics))
        self.assertTrue(isinstance(res_subgroup_feature_has.result, CombinedTestStatistics))
        self.assertLess(res_subgroup_feature_non.result.corrected_test_statistics.statistical_power,
                        res_subgroup_feature_non.result.original_test_statistics.statistical_power)
        self.assertLess(res_subgroup_feature_has.result.corrected_test_statistics.statistical_power,
                        res_subgroup_feature_has.result.original_test_statistics.statistical_power)

    def test_analyze_statistical_test_suite_with_zero_std(self):
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_one_test_zero_std)
        self.assertEqual(res.correction_method, CorrectionMethod.NONE)
        self.assertEqual(len(res.results), 0)

    def test_analyze_statistical_test_with_nan_data(self):
        res = self.getExperiment().analyze_statistical_test_suite(self.suite_with_one_test_with_nan_data)
        self.assertEqual(res.correction_method, CorrectionMethod.NONE)
        self.assertEqual(len(res.results), 0)


class OutlierFilteringTestCases(ExperimentTestCase):
    """ Test outlier filtering and quantile filtering. """

    def test_quantile_filtering_multiple_columns(self):
        exp = self.getExperiment()
        flags = exp._quantile_filtering(
            self.data,
            kpis=[
                'normal_same',
                'normal_shifted',
                'normal_shifted_by_feature',
                'normal_unequal_variance'
            ],
            percentile=99.0,
            threshold_type='upper'
        )
        self.assertEqual(len(flags[flags==True]), 386)
        filtered = [item[1] for item in list(zip(flags, self.data['variant'])) if item[0] == True]
        filtered_dict = dict((val, filtered.count(val)) for val in set(filtered))
        self.assertEqual(filtered_dict['A'], 88)
        self.assertEqual(filtered_dict['B'], 298)

    def test_outlier_filtering_lower_threshold(self):
        exp = self.getExperiment()
        data = exp.outlier_filter(
            self.data,
            kpis=[
                KPI('normal_same'),
                KPI('normal_shifted'),
                KPI('normal_shifted_by_feature'),
                KPI('normal_unequal_variance')
            ],
            percentile=0.1,
            threshold_type='lower'
        )
        self.assertEqual(len(self.data) - len(data), exp.metadata['filtered_entities_number'])
        self.assertEqual(exp.metadata['filtered_entities_per_variant']['A'], 22)
        self.assertEqual(exp.metadata['filtered_entities_per_variant']['B'], 18)

    def test_outlier_filtering_unsupported_kpi(self):
        exp = self.getExperiment()
        with self.assertRaises(KeyError):
            exp.outlier_filter(self.data, kpis=[KPI('revenue')])

    def test_outlier_filtering_unsupported_percentile(self):
        exp = self.getExperiment()
        with self.assertRaises(ValueError):
            exp.outlier_filter(self.data, kpis=[KPI('normal_same')], percentile=101.0)

    def test_outlier_filtering_unsupported_threshold_kind(self):
        exp = self.getExperiment()
        with self.assertRaises(ValueError):
            exp.outlier_filter(self.data, kpis=[KPI('normal_same')], threshold_type='uppper')

    def test_outlier_filtering_high_filtering_percentage(self):
        exp = self.getExperiment()
        with warnings.catch_warnings(record=True) as w:
            exp.outlier_filter(self.data, kpis=[KPI('normal_same')], percentile=97.9)
            self.assertEqual(len(w), 1)
            self.assertTrue(issubclass(w[-1].category, UserWarning))

    def test_outlier_filtering_derived_kpi(self):
        exp = self.getExperiment()
        data = exp.outlier_filter(
            self.data,
            kpis=[
                KPI('normal_same'),
                KPI('normal_shifted'),
                DerivedKPI('derived_kpi', 'normal_same', 'normal_shifted')
            ],
            percentile=0.1,
            threshold_type='lower'
        )
        self.assertIn('derived_kpi', data.columns)

    def test_chi_square_test_result_and_statistics_same_weights(self):
        exp = self.getExperiment()
        data = ['A'] * 23 + ['B'] * 18 + ['C'] * 17 + ['D'] * 19 + ['E'] * 23
        weights = {'A':  0.2, 'B':  0.2, 'C':  0.2, 'D':  0.2, 'E':  0.2}
        result = exp.chi_square_test_result_and_statistics(data, weights)
        self.assertEqual(result[0], True)
        self.assertAlmostEqual(result[1], 0.8087921354109989)
        self.assertAlmostEqual(result[2], 1.6)
        self.assertEqual(result[3]['A'], 23)
        self.assertEqual(result[3]['B'], 18)
        self.assertEqual(result[3]['C'], 17)
        self.assertEqual(result[3]['D'], 19)
        self.assertEqual(result[3]['E'], 23)

        self.assertTrue(all([val == 20 for val in result[4]]))

    def test_chi_square_test_result_and_statistics_different_weights(self):
        exp = self.getExperiment()
        data = ['A'] * 23 + ['B'] * 18 + ['C'] * 17 + ['D'] * 19 + ['E'] * 23
        weights = {'A':  0.25, 'B':  0.15, 'C':  0.10, 'D':  0.40, 'E':  0.10}
        result = exp.chi_square_test_result_and_statistics(data, weights)
        self.assertEqual(result[0], False)
        self.assertAlmostEqual(result[1], 9.064563321754584e-07)
        self.assertAlmostEqual(result[2], 33.585)
        self.assertEqual(len(result[3]), 5)
        self.assertEqual(len(result[4]), 5)

    def test_chi_square_test_result_and_statistics_2_categories(self):
        exp = self.getExperiment()
        data = ['A'] * 17 + ['B'] * 17
        weights = {'A': 0.5, 'B': 0.5}
        result = exp.chi_square_test_result_and_statistics(data, weights)
        self.assertEqual(result[0], True)
        self.assertAlmostEqual(result[1], 1.0)
        self.assertAlmostEqual(result[2], 0.0)
        self.assertEqual(result[3]['A'], 17)
        self.assertEqual(result[4]['A'], 17)
        self.assertEqual(result[3]['B'], 17)
        self.assertEqual(result[4]['B'], 17)

    def test_chi_square_test_result_and_statistics_NaN_data(self):
        exp = self.getExperiment()
        data = ['A'] * 17 + [np.nan] * 17
        weights = {'A': 0.5, 'B': 0.5}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_counts_less_5_2_categories(self):
        exp = self.getExperiment()
        data = ['A'] * 17 + ['B'] * 2 + ['C'] * 3
        weights = {'A': 0.33, 'B': 0.33, 'C': 0.33}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_counts_less_5_1_category(self):
        exp = self.getExperiment()
        data = ['A'] * 8 + ['B'] * 2 + ['C'] * 14
        weights = {'A': 0.33, 'B': 0.33, 'C': 0.33}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_one_category(self):
        exp = self.getExperiment()
        data = ['A'] * 16
        weights = {'A': 0.5}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_empty_weights(self):
        exp = self.getExperiment()
        data = ['A'] * 16
        weights = {}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_no_categories(self):
        exp = self.getExperiment()
        data = []
        weights = {'A': 0.5}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_insufficient_weights(self):
        exp = self.getExperiment()
        data = ['A'] * 16 + ['B'] * 15
        weights = {'A': 0.5}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_None_inputs(self):
        exp = self.getExperiment()
        data = None
        weights = None
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

    def test_chi_square_test_result_and_statistics_empty_inputs(self):
        exp = self.getExperiment()
        data = []
        weights = {}
        with self.assertRaises(ValueError):
            exp.chi_square_test_result_and_statistics(data, weights)

class HelperMethodsTestCases(ExperimentTestCase):
    """ Test other helper methods. """

    # Test is_valid_for_analysis
    def test_is_valid_for_analysis(self):
        exp = self.getExperiment()
        is_valid = exp._is_valid_for_analysis(self.data, self.test_normal_same)
        self.assertTrue(is_valid)

        is_not_valid = exp._is_valid_for_analysis(self.data, self.test_nonsense_variant)
        self.assertFalse(is_not_valid)


if __name__ == '__main__':
    unittest.main()
