"""
Tests for SARIMAX models

Author: Chad Fulton
License: Simplified-BSD
"""

import os
import warnings

from statsmodels.compat.platform import PLATFORM_WIN

import numpy as np
import pandas as pd
import pytest

from statsmodels.tsa.statespace import sarimax, tools
from statsmodels.tsa import arima_model as arima
from .results import results_sarimax
from statsmodels.tools import add_constant
from numpy.testing import (
    assert_, assert_equal, assert_almost_equal, assert_raises, assert_allclose
)


current_path = os.path.dirname(os.path.abspath(__file__))

realgdp_path = os.path.join('results', 'results_realgdpar_stata.csv')
realgdp_results = pd.read_csv(current_path + os.sep + realgdp_path)

coverage_path = os.path.join('results', 'results_sarimax_coverage.csv')
coverage_results = pd.read_csv(os.path.join(current_path, coverage_path))


class TestSARIMAXStatsmodels(object):
    """
    Test ARIMA model using SARIMAX class against statsmodels ARIMA class

    Notes
    -----

    Standard errors are quite good for the OPG case.
    """
    @classmethod
    def setup_class(cls):
        cls.true = results_sarimax.wpi1_stationary
        endog = cls.true['data']

        cls.model_a = arima.ARIMA(endog, order=(1, 1, 1))
        cls.result_a = cls.model_a.fit(disp=-1)

        cls.model_b = sarimax.SARIMAX(endog, order=(1, 1, 1), trend='c',
                                      simple_differencing=True,
                                      hamilton_representation=True)
        cls.result_b = cls.model_b.fit(disp=-1)

    def test_loglike(self):
        assert_allclose(self.result_b.llf, self.result_a.llf)

    def test_aic(self):
        assert_allclose(self.result_b.aic, self.result_a.aic)

    def test_bic(self):
        assert_allclose(self.result_b.bic, self.result_a.bic)

    def test_hqic(self):
        assert_allclose(self.result_b.hqic, self.result_a.hqic)

    def test_mle(self):
        # ARIMA estimates the mean of the process, whereas SARIMAX estimates
        # the intercept. Convert the mean to intercept to compare
        params_a = self.result_a.params.copy()
        params_a[0] = (1 - params_a[1]) * params_a[0]
        assert_allclose(self.result_b.params[:-1], params_a, atol=5e-5)

    def test_bse(self):
        # Test the complex step approximated BSE values
        cpa = self.result_b._cov_params_approx(approx_complex_step=True)
        bse = cpa.diagonal()**0.5
        assert_allclose(bse[1:-1], self.result_a.bse[1:], atol=1e-5)

    def test_t_test(self):
        import statsmodels.tools._testing as smt
        # to trigger failure, un-comment the following:
        #  self.result_b._cache['pvalues'] += 1
        smt.check_ttest_tvalues(self.result_b)
        smt.check_ftest_pvalues(self.result_b)


class TestRealGDPARStata(object):
    """
    Includes tests of filtered states and standardized forecast errors.

    Notes
    -----
    Could also test the usual things like standard errors, etc. but those are
    well-tested elsewhere.
    """
    @classmethod
    def setup_class(cls):
        dlgdp = np.log(realgdp_results['value']).diff()[1:].values
        cls.model = sarimax.SARIMAX(dlgdp, order=(12, 0, 0), trend='n',
                                    hamilton_representation=True)
        # Estimated by Stata
        params = [
            .40725515, .18782621, -.01514009, -.01027267, -.03642297,
            .11576416, .02573029, -.00766572, .13506498, .08649569, .06942822,
            -.10685783, .00007999607
        ]
        cls.results = cls.model.filter(params)

    def test_filtered_state(self):
        for i in range(12):
            assert_allclose(
                realgdp_results.iloc[1:]['u%d' % (i+1)],
                self.results.filter_results.filtered_state[i],
                atol=1e-6
            )

    def test_standardized_forecasts_error(self):
        assert_allclose(
            realgdp_results.iloc[1:]['rstd'],
            self.results.filter_results.standardized_forecasts_error[0],
            atol=1e-3
        )


class SARIMAXStataTests(object):
    def test_loglike(self):
        assert_almost_equal(
            self.result.llf,
            self.true['loglike'], 4
        )

    def test_aic(self):
        assert_almost_equal(
            self.result.aic,
            self.true['aic'], 3
        )

    def test_bic(self):
        assert_almost_equal(
            self.result.bic,
            self.true['bic'], 3
        )

    def test_hqic(self):
        hqic = (
            -2*self.result.llf +
            2*np.log(np.log(self.result.nobs_effective)) *
            self.result.params.shape[0]
        )
        assert_almost_equal(
            self.result.hqic,
            hqic, 3
        )

    def test_standardized_forecasts_error(self):
        cython_sfe = self.result.standardized_forecasts_error
        self.result._standardized_forecasts_error = None
        python_sfe = self.result.standardized_forecasts_error
        assert_allclose(cython_sfe, python_sfe)


class ARIMA(SARIMAXStataTests):
    """
    ARIMA model

    Stata arima documentation, Example 1
    """
    @classmethod
    def setup_class(cls, true, *args, **kwargs):
        cls.true = true
        endog = true['data']

        kwargs.setdefault('simple_differencing', True)
        kwargs.setdefault('hamilton_representation', True)

        cls.model = sarimax.SARIMAX(endog, order=(1, 1, 1), trend='c',
                                    *args, **kwargs)

        # Stata estimates the mean of the process, whereas SARIMAX estimates
        # the intercept of the process. Get the intercept.
        intercept = (1 - true['params_ar'][0]) * true['params_mean'][0]
        params = np.r_[intercept, true['params_ar'], true['params_ma'],
                       true['params_variance']]

        cls.result = cls.model.filter(params)

    def test_mle(self):
        result = self.model.fit(disp=-1)
        assert_allclose(
            result.params, self.result.params,
            atol=1e-3
        )


class TestARIMAStationary(ARIMA):
    """
    Notes
    -----

    Standard errors are very good for the OPG and complex step approximation
    cases.
    """
    @classmethod
    def setup_class(cls):
        super(TestARIMAStationary, cls).setup_class(
            results_sarimax.wpi1_stationary
        )

    def test_bse(self):
        # test defaults
        assert_equal(self.result.cov_type, 'opg')
        assert_equal(self.result._cov_approx_complex_step, True)
        assert_equal(self.result._cov_approx_centered, False)
        # default covariance type (opg)
        assert_allclose(self.result.bse[1], self.true['se_ar_opg'], atol=1e-7)
        assert_allclose(self.result.bse[2], self.true['se_ma_opg'], atol=1e-7)

    def test_bse_approx(self):
        # complex step
        bse = self.result._cov_params_approx(
            approx_complex_step=True).diagonal()**0.5
        assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-7)
        assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-7)

        # The below tests pass irregularly; they give a sense of the precision
        # available with finite differencing
        # finite difference, non-centered
        # with warnings.catch_warnings():
        #     warnings.simplefilter("ignore")
        #     bse = self.result._cov_params_approx(
        #         approx_complex_step=False).diagonal()**0.5
        #     assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-2)
        #     assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-1)

        #     # finite difference, centered
        #     cpa = self.result._cov_params_approx(
        #         approx_complex_step=False, approx_centered=True)
        #     bse = cpa.diagonal()**0.5
        #     assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-3)
        #     assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-3)

    def test_bse_oim(self):
        # OIM covariance type
        oim_bse = self.result.cov_params_oim.diagonal()**0.5
        assert_allclose(oim_bse[1], self.true['se_ar_oim'], atol=1e-3)
        assert_allclose(oim_bse[2], self.true['se_ma_oim'], atol=1e-2)

    def test_bse_robust(self):
        robust_oim_bse = self.result.cov_params_robust_oim.diagonal()**0.5
        cpra = self.result.cov_params_robust_approx
        robust_approx_bse = cpra.diagonal()**0.5
        true_robust_bse = np.r_[
            self.true['se_ar_robust'], self.true['se_ma_robust']
        ]

        assert_allclose(robust_oim_bse[1:3], true_robust_bse, atol=1e-2)
        assert_allclose(robust_approx_bse[1:3], true_robust_bse, atol=1e-3)


class TestARIMADiffuse(ARIMA):
    """
    Notes
    -----

    Standard errors are very good for the OPG and quite good for the complex
    step approximation cases.
    """
    @classmethod
    def setup_class(cls, **kwargs):
        kwargs['initialization'] = 'approximate_diffuse'
        kwargs['initial_variance'] = (
            results_sarimax.wpi1_diffuse['initial_variance']
        )
        super(TestARIMADiffuse, cls).setup_class(results_sarimax.wpi1_diffuse,
                                                 **kwargs)

    def test_bse(self):
        # test defaults
        assert_equal(self.result.cov_type, 'opg')
        assert_equal(self.result._cov_approx_complex_step, True)
        assert_equal(self.result._cov_approx_centered, False)
        # default covariance type (opg)
        assert_allclose(self.result.bse[1], self.true['se_ar_opg'], atol=1e-7)
        assert_allclose(self.result.bse[2], self.true['se_ma_opg'], atol=1e-7)

    def test_bse_approx(self):
        # complex step
        bse = self.result._cov_params_approx(
            approx_complex_step=True).diagonal()**0.5
        assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-4)
        assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-4)

        # The below tests do not pass
        # with warnings.catch_warnings():
        #     warnings.simplefilter("ignore")

        #     # finite difference, non-centered : failure
        #     bse = self.result._cov_params_approx(
        #         approx_complex_step=False).diagonal()**0.5
        #     assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-4)
        #     assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-4)

        #     # finite difference, centered : failure
        #     cpa = self.result._cov_params_approx(
        #         approx_complex_step=False, approx_centered=True)
        #     bse = cpa.diagonal()**0.5
        #     assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-4)
        #     assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-4)

    def test_bse_oim(self):
        # OIM covariance type
        bse = self.result._cov_params_oim().diagonal()**0.5
        assert_allclose(bse[1], self.true['se_ar_oim'], atol=1e-2)
        assert_allclose(bse[2], self.true['se_ma_oim'], atol=1e-1)


class AdditiveSeasonal(SARIMAXStataTests):
    """
    ARIMA model with additive seasonal effects

    Stata arima documentation, Example 2
    """
    @classmethod
    def setup_class(cls, true, *args, **kwargs):
        cls.true = true
        endog = np.log(true['data'])

        kwargs.setdefault('simple_differencing', True)
        kwargs.setdefault('hamilton_representation', True)

        cls.model = sarimax.SARIMAX(
            endog, order=(1, 1, (1, 0, 0, 1)), trend='c', *args, **kwargs
        )

        # Stata estimates the mean of the process, whereas SARIMAX estimates
        # the intercept of the process. Get the intercept.
        intercept = (1 - true['params_ar'][0]) * true['params_mean'][0]
        params = np.r_[intercept, true['params_ar'], true['params_ma'],
                       true['params_variance']]

        cls.result = cls.model.filter(params)