#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Test all the distributions, making sure they do what they say they do.

Based on the Kolmogorov-Smirnov goodness-of-fit test for continuous
distributions, the TestDistributions will actually test the goodness of
all generators listed in :mod:`pypdevsbbl.extra.rngstreams.distributions`.

See Also:
    :mod:`pypdevsbbl.extra.rngstreams.distributions`
"""

# TODO: Use TestU01 instead?

import unittest
from pypdevsbbl.extra.rngstreams import distributions as dist
from scipy import stats
import numpy as np

class TestDistributions(unittest.TestCase):
    def setUp(self):
        N = 10000
        self.range = [x / (N+1) for x in range(1, N+1)]
        self.ALPHA = 0.05   # The p-value must be under ALPHA

    def test_beta(self):
        for a1 in range(1, 5):
            for a2 in range(1, 15):
                x1 = [dist.beta(i, a1, a2) for i in self.range]
                x2 = stats.beta.pdf(self.range, a1, a2)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_binomial(self):
        for n in range(10, 26, 2):
            for q in range(1, 10):
                args = n, q/10
                x = [dist.binomial(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'binom', args=args)
                self.assertLess(p, self.ALPHA)

    def test_cauchy(self):
        for x0 in range(-5, 15):
            for scale in range(1, 10):
                args = x0, scale
                x1 = [dist.cauchy(i, *args) for i in self.range]
                x2 = stats.cauchy.pdf(self.range, *args)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_chiSquare(self):
        for q in range(1, 10):
            args = q,
            x = [dist.chiSquare(i, *args) for i in self.range]
            D, p = stats.kstest(x, 'chi2', args=args)
            self.assertLess(p, self.ALPHA)

    def test_exponential(self):
        for q in range(1, 100):
            args = q/10,
            x = [dist.exponential(i, *args) for i in self.range]
            D, p = stats.kstest(x, 'expon', args=args)
            self.assertLess(p, self.ALPHA)

    def test_F(self):
        for d1 in range(1, 10):
            for d2 in range(10, 50, 5):
                args = d1, d2
                x1 = [dist.F(i, *args) for i in self.range]
                x2 = stats.f.pdf(self.range, *args)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_frechet(self):
        for a in range(1, 10):
            for m in range(0, 60, 15):
                for s in range(1, 5):
                    args = a, m, s
                    x1 = [dist.frechet(i, *args) for i in self.range]
                    x2 = stats.invweibull.pdf(self.range, *args)
                    D, p = stats.ks_2samp(x1, x2)
                    self.assertLess(p, self.ALPHA)

    def test_gamma(self):
        for a in range(1, 10):
            for b in range(1, 10):
                args = a * np.pi / 2, b * np.pi / 2
                x = [dist.gamma(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'gamma', args=args)
                self.assertLess(p, self.ALPHA)

    def test_geometric(self):
        for p in range(1, 100, 5):
            args = p/100, 0
            x = [dist.geometric(i, args[0]) for i in self.range]
            D, p = stats.kstest(x, 'geom', args=args)
            self.assertLess(p, self.ALPHA)

    def test_gumbel(self):
        for b in range(1, 10):
            for m in range(0, 50, 5):
                args = m, b
                x1 = [dist.gumbel(i, *args) for i in self.range]
                x2 = stats.gumbel_r.pdf(self.range, *args)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_JSB(self):
        for a1 in range(1, 5):
            for a2 in range(5, 10):
                for a in range(1, 20, 4):
                    for b in range(1, 15, 3):
                        x = [dist.JSB(i, a1, a2, a, b) for i in self.range]
                        D, p = stats.kstest(x, 'johnsonsb', args=(a1, a2, a, b))
                        self.assertLess(p, self.ALPHA)

    def test_JSU(self):
        for a1 in range(1, 5):
            for a2 in range(5, 10):
                for g in range(1, 20, 4):
                    for b in range(1, 15, 3):
                        args = a1, a2, g, b
                        x1 = [dist.JSU(i, *args) for i in self.range]
                        x2 = stats.johnsonsu.pdf(self.range, *args)
                        D, p = stats.ks_2samp(x1, x2)
                        self.assertLess(p, self.ALPHA)

    def test_laplace(self):
        for b in range(1, 10):
            for m in range(0, 50, 5):
                args = m, b
                x1 = [dist.laplace(i, *args) for i in self.range]
                x2 = stats.laplace.pdf(self.range, *args)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_logistic(self):
        for s in range(1, 10):
            for l in range(0, 50, 5):
                args = l, s
                x1 = [dist.logistic(i, *args) for i in self.range]
                x2 = stats.logistic.pdf(self.range, *args)
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_loglogistic(self):
        for a in range(1, 20):
            for b in range(1, 20):
                x = [dist.loglogistic(i, a, b) for i in self.range]
                D, p = stats.kstest(x, 'fisk', args=(a, b))
                self.assertLess(p, self.ALPHA)

    def test_lognormal(self):
        # Can only test for mean = 0
        for q in range(1, 20):
            args = 0, q
            x1 = [dist.lognormal(i, *args) for i in self.range]
            x2 = stats.lognorm.pdf(self.range, q)
            D, p = stats.ks_2samp(x1, x2)
            self.assertLess(p, self.ALPHA)

    def test_negbinomial(self):
        for n in range(10, 26, 3):
            for q in range(1, 8):
                args = n, q / 10
                x = [dist.negbinomial(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'nbinom', args=args)
                self.assertLess(p, self.ALPHA)

    def test_normal(self):
        for mu in range(0, 10):
            for s in range(1, 5):
                args = mu, s
                x = [dist.normal(i, *args) for i in self.range]
                k2, p = stats.normaltest(x)
                # Other acception range, seeing as it is a two-sided test
                self.assertGreater(p, self.ALPHA/2)
                self.assertLess(p, 1 - self.ALPHA/2)

    def test_pareto(self):
        for s in range(1, 10):
            for a in range(1, 10):
                args = s, a
                x1 = [dist.pareto(i, *args) for i in self.range]
                x2 = stats.pareto.pdf(self.range, args[0], 1, args[1])
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_poisson(self):
        for q in range(5):
            args = q,
            x = [dist.poisson(i, *args) for i in self.range]
            D, p = stats.kstest(x, 'poisson', args=args)
            self.assertLess(p, self.ALPHA)

    def test_PT5(self):
        for a in range(1, 10):
            for b in range(1, 10):
                args = a, b
                x = [dist.pearson5(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'invgamma', args=(a, 1/b))
                self.assertLess(p, self.ALPHA)

    def test_PT6(self):
        for a1 in range(1, 10):
            for a2 in range(1, 10):
                for b in range(1, 10):
                    args = a1, a2, b
                    x1 = [dist.pearson6(i, *args) for i in self.range]
                    x2 = stats.beta.pdf(self.range, a1, a2)
                    x2 = [b * x / (1 - x) for x in x2]
                    D, p = stats.ks_2samp(x1, x2)
                    self.assertLess(p, self.ALPHA)

    def test_student(self):
        for q in range(1, 10):
            args = q,
            x = [dist.student(i, *args) for i in self.range]
            D, p = stats.kstest(x, 't', args=args)
            self.assertLess(p, self.ALPHA)

    def test_triangular(self):
        for a in range(10):
            for b in range(a+1, 10):
                for c in range(b+1, 10):
                    x = [dist.triangular(i, a, b, c) for i in self.range]
                    D, p = stats.kstest(x, 'triang', args=((b-a)/(c-a), a, c-a))
                    self.assertLess(p, self.ALPHA)

    def test_uniform(self):
        for l in range(1, 10):
            for u in range(l+1, 10):
                args = l, u
                x1 = [dist.uniform(i, *args) for i in self.range]
                x2 = stats.uniform.pdf(self.range, args[0], args[1] - args[0])
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_uniformInt(self):
        for l in range(1, 10):
            for u in range(l+1, 10):
                args = l, u
                x1 = [dist.uniformInt(i, *args) for i in self.range]
                x2 = [np.round(i) for i in stats.uniform.pdf(self.range, args[0], args[1] - args[0])]
                D, p = stats.ks_2samp(x1, x2)
                self.assertLess(p, self.ALPHA)

    def test_wald(self):
        for m in range(1, 10):
            for s in range(1, 5):
                args = m, s
                x = [dist.wald(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'wald', args=args)
                self.assertLess(p, self.ALPHA)

    def test_weibull(self):
        for l in range(1, 5):
            for k in range(10, 50, 5):
                args = l, k/10
                x = [dist.weibull(i, *args) for i in self.range]
                D, p = stats.kstest(x, 'dweibull', args=args)
                self.assertLess(p, self.ALPHA)

    def test_zipf(self):
        # Reduce the item size, seeing as zipf has way to much trouble with big datasets
        N = 1000
        r = [x / (N+1) for x in range(1, N+1)]

        # Do the test
        for q in range(10):
            args = q, N
            x1 = [dist.zipf(i, *args) for i in r]
            x2 = stats.zipf.pmf(r, *args)
            D, p = stats.ks_2samp(x1, x2)
            self.assertLess(p, self.ALPHA)

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