# Copyright (c) 2014, Salesforce.com, Inc.  All rights reserved.
# Copyright (c) 2015, Gamelan Labs, Inc.
# Copyright (c) 2016, Google, Inc.
# Copyright (c) 2016, Gamelan Labs, Inc.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# - Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
# - Neither the name of Salesforce.com nor the names of its contributors
#   may be used to endorse or promote products derived from this
#   software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from __future__ import division
try:
    from itertools import izip as zip
except ImportError:
    pass
from itertools import product
import random
from unittest import skip
from unittest import TestCase

import numpy
import scipy.stats
from numpy import pi
from numpy.testing import rand

from goftests import get_dim
from goftests import multinomial_goodness_of_fit
from goftests import discrete_goodness_of_fit
from goftests import auto_density_goodness_of_fit
from goftests import mixed_density_goodness_of_fit
from goftests import split_discrete_continuous
from goftests import volume_of_sphere
from goftests import chi2sf

NUM_BASE_SAMPLES = 250

NUM_SAMPLES_SCALE = 1000

TEST_FAILURE_RATE = 5e-4


class TestMultinomialGoodnessOfFit(TestCase):

    def test_multinomial_goodness_of_fit(self):
        random.seed(0)
        numpy.random.seed(0)
        for dim in range(2, 20):
            sample_count = int(1e5)
            probs = numpy.random.dirichlet([1] * dim)

            counts = numpy.random.multinomial(sample_count, probs)
            p_good = multinomial_goodness_of_fit(probs, counts, sample_count)
            self.assertGreater(p_good, TEST_FAILURE_RATE)

            unif = [1 / dim] * dim
            unif_counts = numpy.random.multinomial(sample_count, unif)
            p_bad = multinomial_goodness_of_fit(probs, unif_counts,
                                                sample_count)
            self.assertLess(p_bad, TEST_FAILURE_RATE)


class TestVolumeOfSphere(TestCase):

    def test_volume_of_sphere(self):
        for r in [0.1, 1.0, 10.0]:
            self.assertAlmostEqual(volume_of_sphere(1, r), 2 * r)
            self.assertAlmostEqual(volume_of_sphere(2, r), pi * r ** 2)
            self.assertAlmostEqual(volume_of_sphere(3, r), 4 / 3 * pi * r ** 3)


SPLIT_EXAMPLES = [
    (False, False, []),
    (0, 0, []),
    ('abc', 'abc', []),
    (0.0, None, [0.0]),
    ((), (), []),
    ([], (), []),
    ((0, ), (0, ), []),
    ([0], (0, ), []),
    ((0.0, ), (None, ), [0.0]),
    ([0.0], (None, ), [0.0]),
    ([True, 1, 'xyz', 3.14, [None, (), ([2.71],)]],
     (True, 1, 'xyz', None, (None, (), ((None,),))),
     [3.14, 2.71]),
    (numpy.zeros(3), (None, None, None), [0.0, 0.0, 0.0]),
]


class TestSplitDiscreteContinuous(TestCase):

    def test_split_continuous_discrete(self):
        for mixed, discrete, continuous in SPLIT_EXAMPLES:
            d, c = split_discrete_continuous(mixed)
            self.assertEqual(d, discrete)
            self.assertAlmostEqual(c, continuous)


class TestChi2CDF(TestCase):

    def test_chi2cdf(self):
        xlist = numpy.linspace(0, 100, 500)
        slist = numpy.arange(1, 41, 1.5)
        for s, x in product(slist, xlist):
            self.assertAlmostEqual(scipy.stats.chi2.sf(x, s), chi2sf(x, s))


class DistributionTestBase(object):
    """Abstract base class for probability distribution unit tests.

    This class supplies two test methods, :meth:`.test_goodness_of_fit`
    and :meth:`.test_mixed_density_goodness_of_fit` for testing the
    goodness of fit functions.

    Subclasses must override and implement one class attribute and two
    instance methods. The :attr:`.dist` class attribute must be set to
    one of SciPy probability distribution constructors in
    :mod:`scipy.stats`. The :meth:`.goodness_of_fit` method must return
    the result of calling one of the goodness of fit functions being
    tested. The :meth:`.probabilites` method must return an object
    representing the probabilities for each sample; the output depends
    on the format of the inputs to the :meth:`.goodness_of_fit` method.

    Subclasses may also set the :attr:`.params` attribute, which is a
    list of tuples that will be provided as arguments to the underlying
    SciPy distribution constructor as specified in :attr:`.dist`. If not
    specified, random arguments will be provided.

    If samples drawn from :attr:`.dist` must be modified in some way
    before the PDF or PMF can be computed, then subclasses may override
    the :meth:`._sample_postprocessing` method.

    """

    #: The SciPy distribution constructor to test.
    dist = None

    #: An optional list of arguments to the distribution constructor.
    #:
    #: Each tuple in this list will be provided as the positional
    #: arguments to the distribution constructor specified in
    #: :attr:`.dist`. If not specified, random arguments will be
    #: provided.
    params = None

    def setUp(self):
        random.seed(0)
        numpy.random.seed(0)

    def _sample_postprocessing(self, sample):
        """Modify a sample drawn from the distribution.

        This method returns a modified version of `sample`, but that
        modification may be arbitrary. This modified sample is the one
        for which the PDF and the goodness-of-fit are computed.

        By default, this is a no-op, but subclasses may wish to override
        this method to modify sample in some way.

        """
        return sample

    def dist_params(self):
        # If there are no parameters, then we provide a random one.
        if self.params is None:
            params = [tuple(1 + rand(self.dist.numargs))]
        else:
            params = self.params
        return params

    def test_mixed_density_goodness_of_fit(self):
        for param in self.dist_params():
            dim = get_dim(self.dist.rvs(*param, size=2)[0])
            sample_count = NUM_BASE_SAMPLES + NUM_SAMPLES_SCALE * dim
            samples = self.dist.rvs(*param, size=sample_count)
            samples = list(map(self._sample_postprocessing, samples))
            probabilities = [self.pdf(sample, *param) for sample in samples]
            gof = mixed_density_goodness_of_fit(samples, probabilities)
            self.assertGreater(gof, TEST_FAILURE_RATE)

    def test_good_fit(self):
        for param in self.dist_params():
            dim = get_dim(self.dist.rvs(*param, size=2)[0])
            sample_count = NUM_BASE_SAMPLES + NUM_SAMPLES_SCALE * dim
            samples = self.dist.rvs(*param, size=sample_count)
            samples = list(map(self._sample_postprocessing, samples))
            probabilities = [self.pdf(sample, *param) for sample in samples]
            gof = self.goodness_of_fit(samples, probabilities)
            self.assertGreater(gof, TEST_FAILURE_RATE)

    def goodness_of_fit(self, samples, probabilities):
        raise NotImplementedError


class ContinuousTestBase(DistributionTestBase):
    """Abstract base class for testing continuous probability distributions.

    Concrete subclasses must set the :attr:`.dist` attribute to be the
    constructor for a continuous probability distribution.

    """

    def goodness_of_fit(self, samples, probabilities):
        gof = auto_density_goodness_of_fit(samples, probabilities)
        return gof

    def pdf(self, *args, **kw):
        return self.dist.pdf(*args, **kw)


class DiscreteTestBase(DistributionTestBase):
    """Abstract base class for testing discrete probability distributions.

    Concrete subclasses must set the :attr:`.dist` attribute to be the
    constructor for a discrete probability distribution.

    """

    def goodness_of_fit(self, samples, probabilities):
        probs_dict = dict(zip(samples, probabilities))
        gof = discrete_goodness_of_fit(samples, probs_dict)
        return gof

    def pdf(self, *args, **kw):
        return self.dist.pmf(*args, **kw)


#
# Multivariate probability distributions.
#

class TestMultivariateNormal(ContinuousTestBase, TestCase):

    dist = scipy.stats.multivariate_normal

    params = [
        (numpy.ones(1), numpy.eye(1)),
        (numpy.ones(2), numpy.eye(2)),
        (numpy.ones(3), numpy.eye(3)),
    ]


class TestDirichlet(ContinuousTestBase, TestCase):

    dist = scipy.stats.dirichlet

    params = [
        ([2.0, 2.5],),
        ([2.0, 2.5, 3.0],),
        ([2.0, 2.5, 3.0, 3.5],),
    ]

    def _sample_postprocessing(self, value):
        """Project onto all but the last dimension."""
        return value[:-1]


#
# Discrete probability distributions.
#

class TestBernoulli(DiscreteTestBase, TestCase):

    dist = scipy.stats.bernoulli

    params = [(0.2, )]


class TestBinomial(DiscreteTestBase, TestCase):

    dist = scipy.stats.binom

    params = [(40, 0.4)]


@skip('')
class TestBoltzmann(DiscreteTestBase, TestCase):

    dist = scipy.stats.boltzmann


class TestDiscreteLaplacian(DiscreteTestBase, TestCase):

    dist = scipy.stats.dlaplace

    params = [(0.8, )]


class TestGeometric(DiscreteTestBase, TestCase):

    dist = scipy.stats.geom

    params = [(0.1, )]