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

from statsmodels.tsa.arima.datasets.brockwell_davis_2002 import dowj, lake
from statsmodels.tsa.arima.estimators.burg import burg
from statsmodels.tsa.innovations.arma_innovations import arma_innovations


@pytest.mark.low_precision(
    "Test against Example 5.1.3 in Brockwell and Davis (2016)"
)
def test_brockwell_davis_example_513():
    # Test against Example 5.1.3 in Brockwell and Davis (2016)
    # (low-precision test, since we are testing against values printed in the
    # textbook)

    # Difference and demean the series
    endog = dowj.diff().iloc[1:]

    # Burg
    res, _ = burg(endog, ar_order=1, demean=True)
    assert_allclose(res.ar_params, [0.4371], atol=1e-4)
    assert_allclose(res.sigma2, 0.1423, atol=1e-4)


@pytest.mark.low_precision(
    "Test against Example 5.1.4 in Brockwell and Davis (2016)"
)
def test_brockwell_davis_example_514():
    # Test against Example 5.1.4 in Brockwell and Davis (2016)
    # (low-precision test, since we are testing against values printed in the
    # textbook)

    # Get the lake data
    endog = lake.copy()

    # Should have 98 observations
    assert_equal(len(endog), 98)
    desired = 9.0041
    assert_allclose(endog.mean(), desired, atol=1e-4)

    # Burg
    res, _ = burg(endog, ar_order=2, demean=True)
    assert_allclose(res.ar_params, [1.0449, -0.2456], atol=1e-4)
    assert_allclose(res.sigma2, 0.4706, atol=1e-4)


def check_itsmr(lake):
    # Test against R itsmr::burg; see results/results_burg.R
    res, _ = burg(lake, 10, demean=True)
    desired_ar_params = [
        1.05853631096,
        -0.32639150878,
        0.04784765122,
        0.02620476111,
        0.04444511374,
        -0.04134010262,
        0.02251178970,
        -0.01427524694,
        0.22223486915,
        -0.20935524387,
    ]
    assert_allclose(res.ar_params, desired_ar_params)

    # itsmr always returns the innovations algorithm estimate of sigma2,
    # whereas we return Burg's estimate
    u, v = arma_innovations(
        np.array(lake) - np.mean(lake), ar_params=res.ar_params, sigma2=1
    )
    desired_sigma2 = 0.4458956354
    assert_allclose(np.sum(u**2 / v) / len(u), desired_sigma2)


def test_itsmr():
    # Note: apparently itsmr automatically demeans (there is no option to
    # control this)
    endog = lake.copy()

    check_itsmr(endog)  # Pandas series
    check_itsmr(endog.values)  # Numpy array
    check_itsmr(endog.tolist())  # Python list


def test_nonstationary_series():
    # Test against R stats::ar.burg; see results/results_burg.R
    endog = np.arange(1, 12) * 1.0
    res, _ = burg(endog, 2, demean=False)

    desired_ar_params = [1.9669331547, -0.9892846679]
    assert_allclose(res.ar_params, desired_ar_params)
    desired_sigma2 = 0.02143066427
    assert_allclose(res.sigma2, desired_sigma2)

    # With var.method = 1, stats::ar.burg also returns something equivalent to
    # the innovations algorithm estimate of sigma2
    u, v = arma_innovations(endog, ar_params=res.ar_params, sigma2=1)
    desired_sigma2 = 0.02191056906
    assert_allclose(np.sum(u**2 / v) / len(u), desired_sigma2)


def test_invalid():
    endog = np.arange(2) * 1.0
    with pytest.raises(ValueError):
        burg(endog, ar_order=2)
    with pytest.raises(ValueError):
        burg(endog, ar_order=-1)
    with pytest.raises(ValueError):
        burg(endog, ar_order=1.5)

    endog = np.arange(10) * 1.0
    with pytest.raises(ValueError):
        burg(endog, ar_order=[1, 3])


def test_misc():
    # Test defaults (order = 0, demean=True)
    endog = lake.copy()
    res, _ = burg(endog)
    assert_allclose(res.params, np.var(endog))

    # Test that integer input gives the same result as float-coerced input.
    endog = np.array([1, 2, 5, 3, -2, 1, -3, 5, 2, 3, -1], dtype=int)
    res_int, _ = burg(endog, 2)
    res_float, _ = burg(endog * 1.0, 2)
    assert_allclose(res_int.params, res_float.params)
