"""

Created on Fri Aug 16 13:41:12 2013

Author: Josef Perktold
"""

import numpy as np
from numpy.testing import assert_allclose, assert_equal
import pytest

# TODO: scipy trim1 is not compatible anymore with my old unit tests
from scipy.stats import trim1

import statsmodels.stats.oneway as smo
from statsmodels.stats.robust_compare import TrimmedMean, trim_mean, trimboth
from statsmodels.tools.testing import Holder


class Test_Trim:
    # test trim functions
    # taken from scipy and adjusted
    def t_est_trim1(self):
        a = np.arange(11)
        assert_equal(trim1(a, 0.1), np.arange(10))
        assert_equal(trim1(a, 0.2), np.arange(9))
        assert_equal(trim1(a, 0.2, tail="left"), np.arange(2, 11))
        assert_equal(trim1(a, 3 / 11.0, tail="left"), np.arange(3, 11))

    def test_trimboth(self):
        a = np.arange(11)
        a2 = np.arange(24).reshape(6, 4)
        a3 = np.arange(24).reshape(6, 4, order="F")
        assert_equal(trimboth(a, 3 / 11.0), np.arange(3, 8))
        assert_equal(trimboth(a, 0.2), np.array([2, 3, 4, 5, 6, 7, 8]))

        assert_equal(trimboth(a2, 0.2), np.arange(4, 20).reshape(4, 4))
        assert_equal(trimboth(a3, 2 / 6.0), np.array([[2, 8, 14, 20], [3, 9, 15, 21]]))
        with pytest.raises(ValueError):
            trimboth(np.arange(24).reshape(4, 6).T, 4 / 6.0)

    def test_trim_mean(self):
        # a = np.array([4, 8, 2, 0, 9, 5, 10, 1, 7, 3, 6])
        idx = np.array([3, 5, 0, 1, 2, 4])
        a2 = np.arange(24).reshape(6, 4)[idx, :]
        a3 = np.arange(24).reshape(6, 4, order="F")[idx, :]
        assert_equal(trim_mean(a3, 2 / 6.0), np.array([2.5, 8.5, 14.5, 20.5]))
        assert_equal(trim_mean(a2, 2 / 6.0), np.array([10.0, 11.0, 12.0, 13.0]))
        idx4 = np.array([1, 0, 3, 2])
        a4 = np.arange(24).reshape(4, 6)[idx4, :]
        assert_equal(
            trim_mean(a4, 2 / 6.0), np.array([9.0, 10.0, 11.0, 12.0, 13.0, 14.0])
        )
        # shuffled arange(24)
        a = np.array(
            [
                7,
                11,
                12,
                21,
                16,
                6,
                22,
                1,
                5,
                0,
                18,
                10,
                17,
                9,
                19,
                15,
                23,
                20,
                2,
                14,
                4,
                13,
                8,
                3,
            ]
        )
        assert_equal(trim_mean(a, 2 / 6.0), 11.5)
        assert_equal(trim_mean([5, 4, 3, 1, 2, 0], 2 / 6.0), 2.5)

        # check axis argument
        np.random.seed(1234)
        a = np.random.randint(20, size=(5, 6, 4, 7))
        for axis in [0, 1, 2, 3, -1]:
            res1 = trim_mean(a, 2 / 6.0, axis=axis)
            res2 = trim_mean(np.rollaxis(a, axis), 2 / 6.0)
            assert_equal(res1, res2)

        res1 = trim_mean(a, 2 / 6.0, axis=None)
        res2 = trim_mean(a.ravel(), 2 / 6.0)
        assert_equal(res1, res2)


class TestTrimmedR1:

    @classmethod
    def setup_class(cls):
        x = np.array(
            [
                77,
                87,
                88,
                114,
                151,
                210,
                219,
                246,
                253,
                262,
                296,
                299,
                306,
                376,
                428,
                515,
                666,
                1310,
                2611,
            ]
        )

        cls.get_results()  # attach k and results
        cls.tm = TrimmedMean(x, cls.k / 19)

    @classmethod
    def get_results(cls):
        cls.k = 1
        # results from R WRS2
        cls.res_basic = np.array(
            [
                342.705882352941,
                92.3342348150314,
                380.157894736842,
                92.9416968861829,
                129679.029239766,
            ]
        )

        # results from R PairedData
        ytt1 = Holder()
        ytt1.statistic = 3.71157981694944
        ytt1.parameter = 16
        ytt1.p_value = 0.00189544440273015
        ytt1.conf_int = np.array([146.966048669017, 538.445716036866])
        ytt1.estimate = 342.705882352941
        ytt1.null_value = 0
        ytt1.alternative = "two.sided"
        ytt1.method = "One sample Yuen test, trim=0.0526315789473684"
        ytt1.data_name = "x"
        cls.ytt1 = ytt1

    def test_basic(self):
        tm = self.tm
        assert_equal(tm.nobs, 19)
        assert_equal(tm.nobs_reduced, 17)
        assert_equal(tm.fraction, self.k / 19)
        assert_equal(tm.data_trimmed.shape[0], tm.nobs_reduced)

        res = [
            tm.mean_trimmed,
            tm.std_mean_trimmed,
            tm.mean_winsorized,
            tm.std_mean_winsorized,
            tm.var_winsorized,
        ]
        assert_allclose(res, self.res_basic, rtol=1e-15)

    def test_inference(self):
        ytt1 = self.ytt1
        tm = self.tm

        ttt = tm.ttest_mean()
        assert_allclose(ttt[0], ytt1.statistic, rtol=1e-13)
        assert_allclose(ttt[1], ytt1.p_value, rtol=1e-13)
        assert_equal(ttt[2], ytt1.parameter)
        assert_allclose(tm.mean_trimmed, ytt1.estimate, rtol=1e-13)

        # regression test for Winsorized t-test,
        # mean, std for it are separately unit tested,
        # df is nobs_reduced-1 in references
        ttw_statistic, ttw_pvalue, tt_w_df = (
            4.090283559190728,
            0.0008537789444194812,
            16,
        )
        ttw = tm.ttest_mean(transform="winsorized")
        assert_allclose(ttw[0], ttw_statistic, rtol=1e-13)
        assert_allclose(ttw[1], ttw_pvalue, rtol=1e-13)
        assert_equal(ttw[2], tt_w_df)

    def test_other(self):
        tm = self.tm
        tm2 = tm.reset_fraction(0.0)
        assert_equal(tm2.nobs_reduced, tm2.nobs)

    @pytest.mark.parametrize("axis", [0, 1])
    def test_vectorized(self, axis):
        tm = self.tm

        x = tm.data
        x2 = np.column_stack((x, 2 * x))
        if axis == 0:
            tm2d = TrimmedMean(x2, self.k / 19, axis=0)
        else:
            tm2d = TrimmedMean(x2.T, self.k / 19, axis=1)
        t1 = [tm.mean_trimmed, 2 * tm.mean_trimmed]
        assert_allclose(tm2d.mean_trimmed, t1, rtol=1e-13)

        t1 = [tm.var_winsorized, 4 * tm.var_winsorized]
        assert_allclose(tm2d.var_winsorized, t1, rtol=1e-13)

        t1 = [tm.std_mean_trimmed, 2 * tm.std_mean_trimmed]
        assert_allclose(tm2d.std_mean_trimmed, t1, rtol=1e-13)

        t1 = [tm.mean_winsorized, 2 * tm.mean_winsorized]
        assert_allclose(tm2d.mean_winsorized, t1, rtol=1e-13)

        t1 = [tm.std_mean_winsorized, 2 * tm.std_mean_winsorized]
        assert_allclose(tm2d.std_mean_winsorized, t1, rtol=1e-13)

        s2, pv2, df2 = tm2d.ttest_mean()
        s, pv, df = tm.ttest_mean()
        assert_allclose(s2, [s, s], rtol=1e-13)
        assert_allclose(pv2, [pv, pv], rtol=1e-13)
        assert_allclose(df2, df, rtol=1e-13)

        s2, pv2, df2 = tm2d.ttest_mean(transform="winsorized")
        s, pv, df = tm.ttest_mean(transform="winsorized")
        assert_allclose(s2, [s, s], rtol=1e-13)
        assert_allclose(pv2, [pv, pv], rtol=1e-13)
        assert_allclose(df2, df, rtol=1e-13)


class TestTrimmedRAnova:

    @classmethod
    def setup_class(cls):
        x = [
            np.array(
                [
                    452.0,
                    874.0,
                    554.0,
                    447.0,
                    356.0,
                    754.0,
                    558.0,
                    574.0,
                    664.0,
                    682.0,
                    547.0,
                    435.0,
                    245.0,
                ]
            ),
            np.array(
                [
                    546.0,
                    547.0,
                    774.0,
                    465.0,
                    459.0,
                    665.0,
                    467.0,
                    365.0,
                    589.0,
                    534.0,
                    456.0,
                    651.0,
                    654.0,
                    665.0,
                    546.0,
                    537.0,
                ]
            ),
            np.array(
                [
                    785.0,
                    458.0,
                    886.0,
                    536.0,
                    669.0,
                    857.0,
                    821.0,
                    772.0,
                    732.0,
                    689.0,
                    654.0,
                    597.0,
                    830.0,
                    827.0,
                ]
            ),
        ]

        cls.x = x
        cls.get_results()  # attach k and results

    @classmethod
    def get_results(cls):
        cls.res_m = [549.3846153846154, 557.5, 722.3571428571429]
        # results from R WRS2
        # > t1w = t1way(y ~ g, df3, tr=1/13)
        cls.res_oneway = Holder(
            test=8.81531710400927,
            df1=2,
            df2=19.8903710685394,
            p_value=0.00181464966984701,
            effsize=0.647137153056774,
        )

        # > yt = yuen(y ~ g, df3[1:29, ], tr=1/13)  # WRS2
        cls.res_2s = Holder(
            test=0.161970203096559,
            conf_int=np.array([-116.437383793431, 99.9568643129114]),
            p_value=0.873436269777141,
            df=15.3931262881751,
            diff=-8.24025974025983,
            effsize=0.0573842557922749,
        )

        # from library onewaytests
        # > bft = bf.test(y ~ g, df3)
        cls.res_bfm = Holder(
            statistic=7.10900606421182,
            parameter=np.array([2, 31.4207256105052]),
            p_value=0.00283841965791224,
            alpha=0.05,
            method="Brown-Forsythe Test",
        )

        # > oww = oneway.test(y ~ g, df3, var.equal = FALSE)
        cls.res_wa = Holder(
            statistic=8.02355212103924,
            parameter=np.array([2, 24.272320628139]),
            p_value=0.00211423625518082,
            method=("One-way analysis of means (not assuming equal variances)"),
        )

        # > ow = oneway.test(y ~ g, df3, var.equal = TRUE)
        cls.res_fa = Holder(
            statistic=7.47403193349076,
            parameter=np.array([2, 40]),
            p_value=0.00174643304119871,
            method="One-way analysis of means",
        )

    def test_oneway(self):
        r1 = self.res_oneway
        r2s = self.res_2s
        res_bfm = self.res_bfm
        res_wa = self.res_wa
        res_fa = self.res_fa

        # check we got the correct data
        m = [x_i.mean() for x_i in self.x]
        assert_allclose(m, self.res_m, rtol=1e-13)

        # 3 sample case
        resg = smo.anova_oneway(self.x, use_var="unequal", trim_frac=1 / 13)
        # assert_allclose(res.statistic, res_bfm.statistic, rtol=1e-13)
        assert_allclose(resg.pvalue, r1.p_value, rtol=1e-13)
        assert_allclose(resg.df, [r1.df1, r1.df2], rtol=1e-13)  # df

        # 2-sample against yuen t-test
        resg = smo.anova_oneway(self.x[:2], use_var="unequal", trim_frac=1 / 13)
        # assert_allclose(res.statistic, res_bfm.statistic, rtol=1e-13)
        assert_allclose(resg.pvalue, r2s.p_value, rtol=1e-13)
        assert_allclose(resg.df, [1, r2s.df], rtol=1e-13)  # df

        # no trimming bfm
        res = smo.anova_oneway(self.x, use_var="bf")
        assert_allclose(res[0], res_bfm.statistic, rtol=1e-13)
        assert_allclose(res.pvalue2, res_bfm.p_value, rtol=1e-13)
        assert_allclose(res.df2, res_bfm.parameter, rtol=1e-13)  # df

        # no trimming welch
        res = smo.anova_oneway(self.x, use_var="unequal")
        assert_allclose(res.statistic, res_wa.statistic, rtol=1e-13)
        assert_allclose(res.pvalue, res_wa.p_value, rtol=1e-13)
        assert_allclose(res.df, res_wa.parameter, rtol=1e-13)  # df

        # no trimming standard anova
        res = smo.anova_oneway(self.x, use_var="equal")
        assert_allclose(res.statistic, res_fa.statistic, rtol=1e-13)
        assert_allclose(res.pvalue, res_fa.p_value, rtol=1e-13)
        assert_allclose(res.df, res_fa.parameter, rtol=1e-13)  # df
