# -*- coding: utf-8 -*-
"""
Created on Sat Apr 16 15:02:13 2011
@author: Josef Perktold
"""
import numpy as np
from numpy.testing import assert_allclose, assert_almost_equal

from statsmodels.sandbox.distributions.multivariate import (
    mvstdnormcdf,
    mvstdtprob,
)
from statsmodels.sandbox.distributions.mv_normal import MVT, MVNormal


class Test_MVN_MVT_prob(object):
    # test for block integratal, cdf, of multivariate t and normal
    # comparison results from R

    @classmethod
    def setup_class(cls):
        cls.corr_equal = np.asarray(
            [[1.0, 0.5, 0.5], [0.5, 1, 0.5], [0.5, 0.5, 1]]
        )
        cls.a = -1 * np.ones(3)
        cls.b = 3 * np.ones(3)
        cls.df = 4

        corr2 = cls.corr_equal.copy()
        corr2[2, 1] = -0.5
        cls.corr2 = corr2

    def test_mvn_mvt_1(self):
        a, b = self.a, self.b
        df = self.df
        corr_equal = self.corr_equal
        # result from R, mvtnorm with option
        # algorithm = GenzBretz(maxpts = 100000, abseps = 0.000001, releps = 0)
        #     or higher
        probmvt_R = 0.60414  # report, ed error approx. 7.5e-06
        probmvn_R = 0.673970  # reported error approx. 6.4e-07
        assert_almost_equal(probmvt_R, mvstdtprob(a, b, corr_equal, df), 4)
        assert_almost_equal(
            probmvn_R, mvstdnormcdf(a, b, corr_equal, abseps=1e-5), 4
        )

        mvn_high = mvstdnormcdf(a, b, corr_equal, abseps=1e-8, maxpts=10000000)
        assert_almost_equal(probmvn_R, mvn_high, 5)
        # this still barely fails sometimes at 6 why?? error is -7.2627419411830374e-007
        # >>> 0.67396999999999996 - 0.67397072627419408
        # -7.2627419411830374e-007
        # >>> assert_almost_equal(0.67396999999999996, 0.67397072627419408, 6)
        # Fail

    def test_mvn_mvt_2(self):
        a, b = self.a, self.b
        df = self.df
        corr2 = self.corr2

        probmvn_R = 0.6472497  # reported error approx. 7.7e-08
        probmvt_R = 0.5881863  # highest reported error up to approx. 1.99e-06
        assert_almost_equal(probmvt_R, mvstdtprob(a, b, corr2, df), 4)
        assert_almost_equal(
            probmvn_R, mvstdnormcdf(a, b, corr2, abseps=1e-5), 4
        )

    def test_mvn_mvt_3(self):
        a, b = self.a, self.b
        df = self.df
        corr2 = self.corr2

        a2 = a.copy()
        a2[:] = -np.inf
        # using higher precision in R, error approx. 6.866163e-07
        probmvn_R = 0.9961141
        # using higher precision in R, error approx. 1.6e-07
        probmvt_R = 0.9522146
        quadkwds = {"epsabs": 1e-08}
        probmvt = mvstdtprob(a2, b, corr2, df, quadkwds=quadkwds)
        assert_allclose(probmvt_R, probmvt, atol=5e-4)
        probmvn = mvstdnormcdf(a2, b, corr2, maxpts=100000, abseps=1e-5)
        assert_allclose(probmvn_R, probmvn, atol=1e-4)

    def test_mvn_mvt_4(self):
        a, bl = self.a, self.b
        df = self.df
        corr2 = self.corr2

        # from 0 to inf
        # print '0 inf'
        a2 = a.copy()
        a2[:] = -np.inf
        probmvn_R = 0.1666667  # error approx. 6.1e-08
        probmvt_R = 0.1666667  # error approx. 8.2e-08
        assert_almost_equal(
            probmvt_R, mvstdtprob(np.zeros(3), -a2, corr2, df), 4
        )
        assert_almost_equal(
            probmvn_R,
            mvstdnormcdf(np.zeros(3), -a2, corr2, maxpts=100000, abseps=1e-5),
            4,
        )

    def test_mvn_mvt_5(self):
        a, bl = self.a, self.b
        df = self.df
        corr2 = self.corr2

        # unequal integration bounds
        # print "ue"
        a3 = np.array([0.5, -0.5, 0.5])
        probmvn_R = (
            0.06910487  # using higher precision in R, error approx. 3.5e-08
        )
        probmvt_R = (
            0.05797867  # using higher precision in R, error approx. 5.8e-08
        )
        assert_almost_equal(mvstdtprob(a3, a3 + 1, corr2, df), probmvt_R, 4)
        assert_almost_equal(
            probmvn_R,
            mvstdnormcdf(a3, a3 + 1, corr2, maxpts=100000, abseps=1e-5),
            4,
        )


class TestMVDistributions(object):
    # this is not well organized

    @classmethod
    def setup_class(cls):
        covx = np.array([[1.0, 0.5], [0.5, 1.0]])
        mu3 = [-1, 0.0, 2.0]
        cov3 = np.array([[1.0, 0.5, 0.75], [0.5, 1.5, 0.6], [0.75, 0.6, 2.0]])
        cls.mu3 = mu3
        cls.cov3 = cov3

        mvn3 = MVNormal(mu3, cov3)
        mvn3c = MVNormal(np.array([0, 0, 0]), cov3)
        cls.mvn3 = mvn3
        cls.mvn3c = mvn3c

    def test_mvn_pdf(self):
        cov3 = self.cov3
        mvn3 = self.mvn3

        r_val = [-7.667977543898155, -6.917977543898155, -5.167977543898155]
        assert_allclose(mvn3.logpdf(cov3), r_val, rtol=1e-13)

        r_val = [
            0.000467562492721686,
            0.000989829804859273,
            0.005696077243833402,
        ]
        assert_allclose(mvn3.pdf(cov3), r_val, rtol=1e-13)

        mvn3b = MVNormal(np.array([0, 0, 0]), cov3)
        r_val = [0.02914269740502042, 0.02269635555984291, 0.01767593948287269]
        assert_allclose(mvn3b.pdf(cov3), r_val, rtol=1e-13)

    def test_mvt_pdf(self, reset_randomstate):
        cov3 = self.cov3
        mu3 = self.mu3

        mvt = MVT((0, 0), 1, 5)
        assert_almost_equal(
            mvt.logpdf(np.array([0.0, 0.0])), -1.837877066409345, decimal=15
        )
        assert_almost_equal(
            mvt.pdf(np.array([0.0, 0.0])), 0.1591549430918953, decimal=15
        )

        mvt.logpdf(np.array([1.0, 1.0])) - (-3.01552989458359)

        mvt1 = MVT((0, 0), 1, 1)
        mvt1.logpdf(np.array([1.0, 1.0])) - (-3.48579549941151)  # decimal=16

        rvs = mvt.rvs(100000)
        assert_almost_equal(np.cov(rvs, rowvar=False), mvt.cov, decimal=1)

        mvt31 = MVT(mu3, cov3, 1)
        assert_almost_equal(
            mvt31.pdf(cov3),
            [
                0.0007276818698165781,
                0.0009980625182293658,
                0.0027661422056214652,
            ],
            decimal=17,
        )

        mvt = MVT(mu3, cov3, 3)
        assert_almost_equal(
            mvt.pdf(cov3),
            [0.000863777424247410, 0.001277510788307594, 0.004156314279452241],
            decimal=17,
        )


if __name__ == "__main__":
    import pytest

    pytest.main([__file__, "-vvs", "-x", "--pdb"])
