# -*- coding: utf-8 -*-
# stationarity tests

from pmdarima.arima.stationarity import ADFTest, PPTest, KPSSTest
from pmdarima.arima.utils import ndiffs
from pmdarima.utils.array import diff
from pmdarima.datasets import load_austres

from sklearn.utils import check_random_state
from numpy.testing import assert_array_almost_equal, assert_almost_equal, \
    assert_array_equal

import numpy as np
import pytest

# for testing rand of len 400 for m==365
random_state = check_random_state(42)
austres = load_austres()


def test_ndiffs_stationary():
    # show that for a stationary vector, ndiffs returns 0
    x = np.ones(10)
    assert ndiffs(x, alpha=0.05, test='kpss', max_d=2) == 0
    assert ndiffs(x, alpha=0.05, test='pp', max_d=2) == 0
    assert ndiffs(x, alpha=0.05, test='adf', max_d=2) == 0


@pytest.mark.parametrize("cls", (KPSSTest, PPTest, ADFTest))
def test_embedding(cls):
    x = np.arange(5)
    expected = np.array([
        [1, 2, 3, 4],
        [0, 1, 2, 3]
    ])

    assert_array_almost_equal(cls._embed(x, 2), expected)

    y = np.array([1, -1, 0, 2, -1, -2, 3])
    assert_array_almost_equal(cls._embed(y, 1),
                              np.array([
                                  [1, -1, 0, 2, -1, -2, 3]
                              ]))

    assert_array_almost_equal(cls._embed(y, 2).T,
                              np.array([
                                  [-1, 1],
                                  [0, -1],
                                  [2, 0],
                                  [-1, 2],
                                  [-2, -1],
                                  [3, -2]
                              ]))

    assert_array_almost_equal(cls._embed(y, 3).T,
                              np.array([
                                  [0, -1, 1],
                                  [2, 0, -1],
                                  [-1, 2, 0],
                                  [-2, -1, 2],
                                  [3, -2, -1]
                              ]))