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

from statsmodels.base.transform import BoxCox
from statsmodels.datasets import macrodata


class TestTransform:

    @classmethod
    def setup_class(cls):
        data = macrodata.load_pandas()
        cls.x = data.data["realgdp"].values
        cls.bc = BoxCox()

    def test_nonpositive(self):
        # Testing negative values
        y = [1, -1, 1]
        with pytest.raises(ValueError):
            self.bc.transform_boxcox(y)

        # Testing nonzero
        y = [1, 0, 1]
        with pytest.raises(ValueError):
            self.bc.transform_boxcox(y)

    def test_invalid_bounds(self):
        # more than two bounds
        with pytest.raises(ValueError):
            self.bc._est_lambda(self.x, (-3, 2, 3))

        # upper bound <= lower bound
        with pytest.raises(ValueError):
            self.bc._est_lambda(self.x, (2, -1))

    def test_unclear_methods(self):
        # Both _est_lambda and untransform have a method argument that should
        # be tested.
        with pytest.raises(ValueError):
            self.bc._est_lambda(self.x, (-1, 2), "test")
        with pytest.raises(ValueError):
            self.bc.untransform_boxcox(self.x, 1, "test")

    def test_unclear_scale_parameter(self):
        # bc.guerrero allows for 'mad' and 'sd', for the MAD and Standard
        # Deviation, respectively
        with pytest.raises(ValueError):
            self.bc._est_lambda(self.x, scale="test")

        # Next, check if mad/sd work:
        self.bc._est_lambda(self.x, scale="mad")
        self.bc._est_lambda(self.x, scale="MAD")

        self.bc._est_lambda(self.x, scale="sd")
        self.bc._est_lambda(self.x, scale="SD")

    def test_valid_guerrero(self):
        # `l <- BoxCox.lambda(x, method="guerrero")` on a ts object
        # with frequency 4 (BoxCox.lambda defaults to 2, but we use
        # Guerrero and Perera (2004) as a guideline)
        lmbda = self.bc._est_lambda(self.x, method="guerrero", window_length=4)
        assert_almost_equal(lmbda, 0.507624, 4)

        # `l <- BoxCox.lambda(x, method="guerrero")` with the default grouping
        # parameter (namely, window_length=2).
        lmbda = self.bc._est_lambda(self.x, method="guerrero", window_length=2)
        assert_almost_equal(lmbda, 0.513893, 4)

    def test_guerrero_robust_scale(self):
        # The lambda is derived from a manual check of the values for the MAD.
        # Compare also the result for the standard deviation on R=4: 0.5076,
        # i.e. almost the same value.
        lmbda = self.bc._est_lambda(self.x, scale="mad")
        assert_almost_equal(lmbda, 0.488621, 4)

    def test_loglik_lambda_estimation(self):
        # 0.2 is the value returned by `BoxCox.lambda(x, method="loglik")`
        lmbda = self.bc._est_lambda(self.x, method="loglik")
        assert_almost_equal(lmbda, 0.2, 1)

    def test_boxcox_transformation_methods(self):
        # testing estimated lambda vs. provided. Should result in almost
        # the same transformed data. Value taken from R.
        y_transformed_no_lambda = self.bc.transform_boxcox(self.x)
        y_transformed_lambda = self.bc.transform_boxcox(self.x, 0.507624)

        assert_almost_equal(y_transformed_no_lambda[0], y_transformed_lambda[0], 3)

        # a perfectly increasing set has a constant variance over the entire
        # series, hence stabilising should result in the same scale: lmbda = 1.
        y, lmbda = self.bc.transform_boxcox(np.arange(1, 100))
        assert_almost_equal(lmbda, 1.0, 5)

    def test_zero_lambda(self):
        # zero lambda should be a log transform.
        y_transform_zero_lambda, lmbda = self.bc.transform_boxcox(self.x, 0.0)

        assert_equal(lmbda, 0.0)
        assert_almost_equal(y_transform_zero_lambda, np.log(self.x), 5)

    def test_naive_back_transformation(self):
        # test both transformations functions -> 0. and .5
        y_zero_lambda = self.bc.transform_boxcox(self.x, 0.0)
        y_half_lambda = self.bc.transform_boxcox(self.x, 0.5)

        y_zero_lambda_un = self.bc.untransform_boxcox(*y_zero_lambda, method="naive")
        y_half_lambda_un = self.bc.untransform_boxcox(*y_half_lambda, method="naive")

        assert_almost_equal(self.x, y_zero_lambda_un, 5)
        assert_almost_equal(self.x, y_half_lambda_un, 5)
