from __future__ import division, absolute_import, print_function

import warnings

import numpy as np

try:
    import scipy.stats as stats
except ImportError:
    pass

from .common import Benchmark

class Anderson_KSamp(Benchmark):
    def setup(self, *args):
        self.rand = [np.random.normal(loc=i, size=1000) for i in range(3)]

    def time_anderson_ksamp(self):
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', UserWarning)
            stats.anderson_ksamp(self.rand)


class CorrelationFunctions(Benchmark):
    param_names = ['alternative']
    params = [
        ['two-sided', 'less', 'greater']
    ]

    def setup(self, mode):
        a = np.random.rand(2,2) * 10
        self.a = a

    def time_fisher_exact(self, alternative):   
        oddsratio, pvalue = stats.fisher_exact(self.a, alternative=alternative)


class InferentialStats(Benchmark):
    def setup(self):
        np.random.seed(12345678)
        self.a = stats.norm.rvs(loc=5, scale=10, size=500)
        self.b = stats.norm.rvs(loc=8, scale=10, size=20)
        self.c = stats.norm.rvs(loc=8, scale=20, size=20)

    def time_ttest_ind_same_var(self):
        # test different sized sample with variances
        stats.ttest_ind(self.a, self.b)
        stats.ttest_ind(self.a, self.b, equal_var=False)

    def time_ttest_ind_diff_var(self):
        # test different sized sample with different variances
        stats.ttest_ind(self.a, self.c)
        stats.ttest_ind(self.a, self.c, equal_var=False)


class Distribution(Benchmark):
    param_names = ['distribution', 'properties']    
    params = [
        ['cauchy', 'gamma', 'beta'],
        ['pdf', 'cdf', 'rvs', 'fit']
    ]
    
    def setup(self, distribution, properties):
        np.random.seed(12345678)
        self.x = np.random.rand(100)

    def time_distribution(self, distribution, properties):
        if distribution == 'gamma':
            if properties == 'pdf':
                stats.gamma.pdf(self.x, a=5, loc=4, scale=10)
            elif properties == 'cdf':
                stats.gamma.cdf(self.x, a=5, loc=4, scale=10)
            elif properties == 'rvs':
                stats.gamma.rvs(size=1000, a=5, loc=4, scale=10)
            elif properties == 'fit':
                stats.gamma.fit(self.x, loc=4, scale=10)
        elif distribution == 'cauchy':
            if properties == 'pdf':
                stats.cauchy.pdf(self.x, loc=4, scale=10)
            elif properties == 'cdf':
                stats.cauchy.cdf(self.x, loc=4, scale=10)
            elif properties == 'rvs':
                stats.cauchy.rvs(size=1000, loc=4, scale=10)
            elif properties == 'fit':
                stats.cauchy.fit(self.x, loc=4, scale=10)
        elif distribution == 'beta':
            if properties == 'pdf':
                stats.beta.pdf(self.x, a=5, b=3, loc=4, scale=10)
            elif properties == 'cdf':
                stats.beta.cdf(self.x, a=5, b=3, loc=4, scale=10)
            elif properties == 'rvs':
                stats.beta.rvs(size=1000, a=5, b=3, loc=4, scale=10)
            elif properties == 'fit':
                stats.beta.fit(self.x, loc=4, scale=10)

    # Retain old benchmark results (remove this if changing the benchmark)
    time_distribution.version = "fb22ae5386501008d945783921fe44aef3f82c1dafc40cddfaccaeec38b792b0"


class DescriptiveStats(Benchmark):
    param_names = ['n_levels']
    params = [
        [10, 1000]
    ]

    def setup(self, n_levels):
        np.random.seed(12345678)
        self.levels = np.random.randint(n_levels, size=(1000, 10))

    def time_mode(self, n_levels):
        stats.mode(self.levels, axis=0)


[(35, 'scipy.stats.fisher_exact', 'stats.fisher_exact', True, 'import scipy.stats as stats\n'), (40, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (41, 'scipy.stats.norm.rvs', 'stats.norm.rvs', True, 'import scipy.stats as stats\n'), (42, 'scipy.stats.norm.rvs', 'stats.norm.rvs', True, 'import scipy.stats as stats\n'), (43, 'scipy.stats.norm.rvs', 'stats.norm.rvs', True, 'import scipy.stats as stats\n'), (47, 'scipy.stats.ttest_ind', 'stats.ttest_ind', True, 'import scipy.stats as stats\n'), (48, 'scipy.stats.ttest_ind', 'stats.ttest_ind', True, 'import scipy.stats as stats\n'), (52, 'scipy.stats.ttest_ind', 'stats.ttest_ind', True, 'import scipy.stats as stats\n'), (53, 'scipy.stats.ttest_ind', 'stats.ttest_ind', True, 'import scipy.stats as stats\n'), (64, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (65, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (107, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (108, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (111, 'scipy.stats.mode', 'stats.mode', True, 'import scipy.stats as stats\n'), (16, 'numpy.random.normal', 'np.random.normal', True, 'import numpy as np\n'), (19, 'warnings.catch_warnings', 'warnings.catch_warnings', False, 'import warnings\n'), (20, 'warnings.simplefilter', 'warnings.simplefilter', False, 'import warnings\n'), (21, 'scipy.stats.anderson_ksamp', 'stats.anderson_ksamp', True, 'import scipy.stats as stats\n'), (31, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (70, 'scipy.stats.gamma.pdf', 'stats.gamma.pdf', True, 'import scipy.stats as stats\n'), (72, 'scipy.stats.gamma.cdf', 'stats.gamma.cdf', True, 'import scipy.stats as stats\n'), (79, 'scipy.stats.cauchy.pdf', 'stats.cauchy.pdf', True, 'import scipy.stats as stats\n'), (74, 'scipy.stats.gamma.rvs', 'stats.gamma.rvs', True, 'import scipy.stats as stats\n'), (81, 'scipy.stats.cauchy.cdf', 'stats.cauchy.cdf', True, 'import scipy.stats as stats\n'), (88, 'scipy.stats.beta.pdf', 'stats.beta.pdf', True, 'import scipy.stats as stats\n'), (76, 'scipy.stats.gamma.fit', 'stats.gamma.fit', True, 'import scipy.stats as stats\n'), (83, 'scipy.stats.cauchy.rvs', 'stats.cauchy.rvs', True, 'import scipy.stats as stats\n'), (90, 'scipy.stats.beta.cdf', 'stats.beta.cdf', True, 'import scipy.stats as stats\n'), (85, 'scipy.stats.cauchy.fit', 'stats.cauchy.fit', True, 'import scipy.stats as stats\n'), (92, 'scipy.stats.beta.rvs', 'stats.beta.rvs', True, 'import scipy.stats as stats\n'), (94, 'scipy.stats.beta.fit', 'stats.beta.fit', True, 'import scipy.stats as stats\n')]