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

from statsmodels.tsa.arima import specification
from statsmodels.tsa.statespace.tools import (
    constrain_stationary_univariate as constrain,
    unconstrain_stationary_univariate as unconstrain,
)


def check_attributes(
    spec,
    order,
    seasonal_order,
    enforce_stationarity,
    enforce_invertibility,
    concentrate_scale,
):
    p, d, q = order
    P, D, Q, s = seasonal_order

    assert_equal(spec.order, (p, d, q))
    assert_equal(spec.seasonal_order, (P, D, Q, s))

    assert_equal(spec.ar_order, p)
    assert_equal(spec.diff, d)
    assert_equal(spec.ma_order, q)

    assert_equal(spec.seasonal_ar_order, P)
    assert_equal(spec.seasonal_diff, D)
    assert_equal(spec.seasonal_ma_order, Q)
    assert_equal(spec.seasonal_periods, s)

    assert_equal(spec.ar_lags, (p if isinstance(p, list) else np.arange(1, p + 1)))
    assert_equal(spec.ma_lags, (q if isinstance(q, list) else np.arange(1, q + 1)))

    assert_equal(
        spec.seasonal_ar_lags, (P if isinstance(P, list) else np.arange(1, P + 1))
    )
    assert_equal(
        spec.seasonal_ma_lags, (Q if isinstance(Q, list) else np.arange(1, Q + 1))
    )

    max_ar_order = p[-1] if isinstance(p, list) else p
    max_ma_order = q[-1] if isinstance(q, list) else q
    max_seasonal_ar_order = P[-1] if isinstance(P, list) else P
    max_seasonal_ma_order = Q[-1] if isinstance(Q, list) else Q
    assert_equal(spec.max_ar_order, max_ar_order)
    assert_equal(spec.max_ma_order, max_ma_order)
    assert_equal(spec.max_seasonal_ar_order, max_seasonal_ar_order)
    assert_equal(spec.max_seasonal_ma_order, max_seasonal_ma_order)
    assert_equal(spec.max_reduced_ar_order, max_ar_order + max_seasonal_ar_order * s)
    assert_equal(spec.max_reduced_ma_order, max_ma_order + max_seasonal_ma_order * s)

    assert_equal(spec.enforce_stationarity, enforce_stationarity)
    assert_equal(spec.enforce_invertibility, enforce_invertibility)
    assert_equal(spec.concentrate_scale, concentrate_scale)


def check_properties(
    spec,
    order,
    seasonal_order,
    enforce_stationarity,
    enforce_invertibility,
    concentrate_scale,
    is_ar_consecutive,
    is_ma_consecutive,
    exog_names,
    ar_names,
    ma_names,
    seasonal_ar_names,
    seasonal_ma_names,
):
    p, d, q = order
    P, D, Q, s = seasonal_order

    k_exog_params = len(exog_names)
    k_ar_params = len(p) if isinstance(p, list) else p
    k_ma_params = len(q) if isinstance(q, list) else q
    k_seasonal_ar_params = len(P) if isinstance(P, list) else P
    k_seasonal_ma_params = len(Q) if isinstance(Q, list) else Q
    k_variance_params = int(not concentrate_scale)

    param_names = (
        exog_names + ar_names + ma_names + seasonal_ar_names + seasonal_ma_names
    )
    if not concentrate_scale:
        param_names.append("sigma2")

    assert_equal(spec.is_ar_consecutive, is_ar_consecutive)
    assert_equal(spec.is_ma_consecutive, is_ma_consecutive)
    assert_equal(spec.is_integrated, d + D > 0)
    assert_equal(spec.is_seasonal, s > 0)

    assert_equal(spec.k_exog_params, k_exog_params)
    assert_equal(spec.k_ar_params, k_ar_params)
    assert_equal(spec.k_ma_params, k_ma_params)
    assert_equal(spec.k_seasonal_ar_params, k_seasonal_ar_params)
    assert_equal(spec.k_seasonal_ma_params, k_seasonal_ma_params)
    assert_equal(
        spec.k_params,
        k_exog_params
        + k_ar_params
        + k_ma_params
        + k_seasonal_ar_params
        + k_seasonal_ma_params
        + k_variance_params,
    )

    assert_equal(spec.exog_names, exog_names)
    assert_equal(spec.ar_names, ar_names)
    assert_equal(spec.ma_names, ma_names)
    assert_equal(spec.seasonal_ar_names, seasonal_ar_names)
    assert_equal(spec.seasonal_ma_names, seasonal_ma_names)
    assert_equal(spec.param_names, param_names)


def check_methods(
    spec,
    order,
    seasonal_order,
    enforce_stationarity,
    enforce_invertibility,
    concentrate_scale,
    exog_params,
    ar_params,
    ma_params,
    seasonal_ar_params,
    seasonal_ma_params,
    sigma2,
):
    params = np.r_[
        exog_params,
        ar_params,
        ma_params,
        seasonal_ar_params,
        seasonal_ma_params,
        sigma2,
    ]

    # Test methods
    desired = {
        "exog_params": exog_params,
        "ar_params": ar_params,
        "ma_params": ma_params,
        "seasonal_ar_params": seasonal_ar_params,
        "seasonal_ma_params": seasonal_ma_params,
    }
    if not concentrate_scale:
        desired["sigma2"] = sigma2
    assert_equal(spec.split_params(params), desired)

    assert_equal(spec.join_params(**desired), params)

    assert_equal(spec.validate_params(params), None)

    # Wrong shape
    with pytest.raises(ValueError):
        spec.validate_params([])

    # Wrong dtype
    with pytest.raises(ValueError):
        spec.validate_params(["a"] + params[1:].tolist())

    # NaN / Infinity
    with pytest.raises(ValueError):
        spec.validate_params(np.r_[np.inf, params[1:]])
    with pytest.raises(ValueError):
        spec.validate_params(np.r_[np.nan, params[1:]])

    # Non-stationary / non-invertible
    if spec.max_ar_order > 0:
        params = np.r_[
            exog_params,
            np.ones_like(ar_params),
            ma_params,
            np.zeros_like(seasonal_ar_params),
            seasonal_ma_params,
            sigma2,
        ]
        if enforce_stationarity:
            with pytest.raises(ValueError):
                spec.validate_params(params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_ma_order > 0:
        params = np.r_[
            exog_params,
            ar_params,
            np.ones_like(ma_params),
            seasonal_ar_params,
            np.zeros_like(seasonal_ma_params),
            sigma2,
        ]
        if enforce_invertibility:
            with pytest.raises(ValueError):
                spec.validate_params(params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_seasonal_ar_order > 0:
        params = np.r_[
            exog_params,
            np.zeros_like(ar_params),
            ma_params,
            np.ones_like(seasonal_ar_params),
            seasonal_ma_params,
            sigma2,
        ]
        if enforce_stationarity:
            with pytest.raises(ValueError):
                spec.validate_params(params)
        else:
            assert_equal(spec.validate_params(params), None)
    if spec.max_seasonal_ma_order > 0:
        params = np.r_[
            exog_params,
            ar_params,
            np.zeros_like(ma_params),
            seasonal_ar_params,
            np.ones_like(seasonal_ma_params),
            sigma2,
        ]
        if enforce_invertibility:
            with pytest.raises(ValueError):
                spec.validate_params(params)
        else:
            assert_equal(spec.validate_params(params), None)

    # Invalid variances
    if not concentrate_scale:
        params = np.r_[
            exog_params,
            ar_params,
            ma_params,
            seasonal_ar_params,
            seasonal_ma_params,
            0.0,
        ]
        with pytest.raises(ValueError):
            spec.validate_params(params)
        params = np.r_[
            exog_params,
            ar_params,
            ma_params,
            seasonal_ar_params,
            seasonal_ma_params,
            -1,
        ]
        with pytest.raises(ValueError):
            spec.validate_params(params)

    # Constrain / unconstrain
    unconstrained_ar_params = ar_params
    unconstrained_ma_params = ma_params
    unconstrained_seasonal_ar_params = seasonal_ar_params
    unconstrained_seasonal_ma_params = seasonal_ma_params
    unconstrained_sigma2 = sigma2

    if spec.max_ar_order > 0 and enforce_stationarity:
        unconstrained_ar_params = unconstrain(np.array(ar_params))
    if spec.max_ma_order > 0 and enforce_invertibility:
        unconstrained_ma_params = unconstrain(-np.array(ma_params))
    if spec.max_seasonal_ar_order > 0 and enforce_stationarity:
        unconstrained_seasonal_ar_params = unconstrain(np.array(seasonal_ar_params))
    if spec.max_seasonal_ma_order > 0 and enforce_invertibility:
        unconstrained_seasonal_ma_params = unconstrain(
            -np.array(unconstrained_seasonal_ma_params)
        )
    if not concentrate_scale:
        unconstrained_sigma2 = unconstrained_sigma2**0.5

    unconstrained_params = np.r_[
        exog_params,
        unconstrained_ar_params,
        unconstrained_ma_params,
        unconstrained_seasonal_ar_params,
        unconstrained_seasonal_ma_params,
        unconstrained_sigma2,
    ]
    params = np.r_[
        exog_params,
        ar_params,
        ma_params,
        seasonal_ar_params,
        seasonal_ma_params,
        sigma2,
    ]

    assert_allclose(spec.unconstrain_params(params), unconstrained_params)

    assert_allclose(spec.constrain_params(unconstrained_params), params)

    assert_allclose(spec.constrain_params(spec.unconstrain_params(params)), params)


@pytest.mark.parametrize(
    "n,d,D,s,params,which",
    [
        # AR models
        (0, 0, 0, 0, np.array([1.0]), "p"),
        (1, 0, 0, 0, np.array([0.5, 1.0]), "p"),
        (1, 0, 0, 0, np.array([-0.2, 100.0]), "p"),
        (2, 0, 0, 0, np.array([-0.2, 0.5, 100.0]), "p"),
        (20, 0, 0, 0, np.array([0.0] * 20 + [100.0]), "p"),
        # ARI models
        (0, 1, 0, 0, np.array([1.0]), "p"),
        (0, 1, 1, 4, np.array([1.0]), "p"),
        (1, 1, 0, 0, np.array([0.5, 1.0]), "p"),
        (1, 1, 1, 4, np.array([0.5, 1.0]), "p"),
        # MA models
        (0, 0, 0, 0, np.array([1.0]), "q"),
        (1, 0, 0, 0, np.array([0.5, 1.0]), "q"),
        (1, 0, 0, 0, np.array([-0.2, 100.0]), "q"),
        (2, 0, 0, 0, np.array([-0.2, 0.5, 100.0]), "q"),
        (20, 0, 0, 0, np.array([0.0] * 20 + [100.0]), "q"),
        # IMA models
        (0, 1, 0, 0, np.array([1.0]), "q"),
        (0, 1, 1, 4, np.array([1.0]), "q"),
        (1, 1, 0, 0, np.array([0.5, 1.0]), "q"),
        (1, 1, 1, 4, np.array([0.5, 1.0]), "q"),
    ],
)
def test_specification_ar_or_ma(n, d, D, s, params, which):
    if which == "p":
        # d unchanged
        p, q = n, 0
        ar_names = ["ar.L%d" % i for i in range(1, p + 1)]
        ma_names = []
    else:
        # d unchanged
        p, q = 0, n
        ar_names = []
        ma_names = ["ma.L%d" % i for i in range(1, q + 1)]
    ar_params = params[:p]
    ma_params = params[p:-1]
    sigma2 = params[-1]
    # D, s unchanged
    P, Q = 0, 0

    args = ((p, d, q), (P, D, Q, s))
    kwargs = {
        "enforce_stationarity": None,
        "enforce_invertibility": None,
        "concentrate_scale": None,
    }

    properties_kwargs = kwargs.copy()
    properties_kwargs.update(
        {
            "is_ar_consecutive": True,
            "is_ma_consecutive": True,
            "exog_names": [],
            "ar_names": ar_names,
            "ma_names": ma_names,
            "seasonal_ar_names": [],
            "seasonal_ma_names": [],
        }
    )

    methods_kwargs = kwargs.copy()
    methods_kwargs.update(
        {
            "exog_params": [],
            "ar_params": ar_params,
            "ma_params": ma_params,
            "seasonal_ar_params": [],
            "seasonal_ma_params": [],
            "sigma2": sigma2,
        }
    )

    # Test the spec created with order, seasonal_order
    spec = specification.SARIMAXSpecification(
        order=(p, d, q), seasonal_order=(P, D, Q, s)
    )

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)

    # Test the spec created with ar_order, etc.
    spec = specification.SARIMAXSpecification(
        ar_order=p,
        diff=d,
        ma_order=q,
        seasonal_ar_order=P,
        seasonal_diff=D,
        seasonal_ma_order=Q,
        seasonal_periods=s,
    )

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)


@pytest.mark.parametrize(
    (
        "endog,exog,p,d,q,P,D,Q,s,"
        "enforce_stationarity,enforce_invertibility,"
        "concentrate_scale"
    ),
    [
        (None, None, 0, 0, 0, 0, 0, 0, 0, True, True, False),
        (None, None, 1, 0, 1, 0, 0, 0, 0, True, True, False),
        (None, None, 1, 1, 1, 0, 0, 0, 0, True, True, False),
        (None, None, 1, 0, 0, 0, 0, 0, 4, True, True, False),
        (None, None, 0, 0, 0, 1, 1, 1, 4, True, True, False),
        (None, None, 1, 0, 0, 1, 0, 0, 4, True, True, False),
        (None, None, 1, 0, 0, 1, 1, 1, 4, True, True, False),
        (None, None, 2, 1, 3, 4, 1, 3, 12, True, True, False),
        # Non-consecutive lag orders
        (None, None, [1, 3], 0, 0, 1, 0, 0, 4, True, True, False),
        (None, None, 0, 0, 0, 0, 0, [1, 3], 4, True, True, False),
        (None, None, [2], 0, [1, 3], [1, 3], 0, [1, 4], 4, True, True, False),
        # Modify enforce / concentrate
        (None, None, 2, 1, 3, 4, 1, 3, 12, False, False, True),
        (None, None, 2, 1, 3, 4, 1, 3, 12, True, False, True),
        (None, None, 2, 1, 3, 4, 1, 3, 12, False, True, True),
        # Endog / exog
        (True, None, 2, 1, 3, 4, 1, 3, 12, False, True, True),
        (None, 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
        (True, 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
        ("y", None, 2, 1, 3, 4, 1, 3, 12, False, True, True),
        (None, ["x1"], 2, 1, 3, 4, 1, 3, 12, False, True, True),
        ("y", ["x1"], 2, 1, 3, 4, 1, 3, 12, False, True, True),
        ("y", ["x1", "x2"], 2, 1, 3, 4, 1, 3, 12, False, True, True),
        (True, ["x1", "x2"], 2, 1, 3, 4, 1, 3, 12, False, True, True),
        ("y", 2, 2, 1, 3, 4, 1, 3, 12, False, True, True),
    ],
)
def test_specification(
    endog,
    exog,
    p,
    d,
    q,
    P,
    D,
    Q,
    s,
    enforce_stationarity,
    enforce_invertibility,
    concentrate_scale,
):
    # Assumptions:
    # - p, q, P, Q are either integers or lists of non-consecutive integers
    #   (i.e. we are not testing boolean lists or consecutive lists here, which
    #   should be tested in the `standardize_lag_order` tests)

    # Construct the specification
    if isinstance(p, list):
        k_ar_params = len(p)
        max_ar_order = p[-1]
    else:
        k_ar_params = max_ar_order = p

    if isinstance(q, list):
        k_ma_params = len(q)
        max_ma_order = q[-1]
    else:
        k_ma_params = max_ma_order = q

    if isinstance(P, list):
        k_seasonal_ar_params = len(P)
        max_seasonal_ar_order = P[-1]
    else:
        k_seasonal_ar_params = max_seasonal_ar_order = P

    if isinstance(Q, list):
        k_seasonal_ma_params = len(Q)
        max_seasonal_ma_order = Q[-1]
    else:
        k_seasonal_ma_params = max_seasonal_ma_order = Q

    # Get endog / exog
    nobs = (
        d
        + D * s
        + max(
            3 * max_ma_order + 1,
            3 * max_seasonal_ma_order * s + 1,
            max_ar_order,
            max_seasonal_ar_order * s,
        )
        + 1
    )

    if endog is True:
        endog = np.arange(nobs) * 1.0
    elif isinstance(endog, str):
        endog = pd.Series(np.arange(nobs) * 1.0, name=endog)
    elif endog is not None:
        raise ValueError("Invalid `endog` in test setup.")

    if isinstance(exog, int):
        exog_names = ["x%d" % (i + 1) for i in range(exog)]
        exog = np.arange(nobs * len(exog_names)).reshape(nobs, len(exog_names))
    elif isinstance(exog, list):
        exog_names = exog
        exog = np.arange(nobs * len(exog_names)).reshape(nobs, len(exog_names))
        exog = pd.DataFrame(exog, columns=exog_names)
    elif exog is None:
        exog_names = []
    else:
        raise ValueError("Invalid `exog` in test setup.")

    # Setup args, kwargs
    args = ((p, d, q), (P, D, Q, s))
    kwargs = {
        "enforce_stationarity": enforce_stationarity,
        "enforce_invertibility": enforce_invertibility,
        "concentrate_scale": concentrate_scale,
    }
    properties_kwargs = kwargs.copy()
    is_ar_consecutive = not isinstance(p, list) and max_seasonal_ar_order == 0
    is_ma_consecutive = not isinstance(q, list) and max_seasonal_ma_order == 0
    properties_kwargs.update(
        {
            "is_ar_consecutive": is_ar_consecutive,
            "is_ma_consecutive": is_ma_consecutive,
            "exog_names": exog_names,
            "ar_names": [
                "ar.L%d" % i for i in (p if isinstance(p, list) else range(1, p + 1))
            ],
            "ma_names": [
                "ma.L%d" % i for i in (q if isinstance(q, list) else range(1, q + 1))
            ],
            "seasonal_ar_names": [
                "ar.S.L%d" % (i * s)
                for i in (P if isinstance(P, list) else range(1, P + 1))
            ],
            "seasonal_ma_names": [
                "ma.S.L%d" % (i * s)
                for i in (Q if isinstance(Q, list) else range(1, Q + 1))
            ],
        }
    )

    methods_kwargs = kwargs.copy()
    methods_kwargs.update(
        {
            "exog_params": np.arange(len(exog_names)),
            "ar_params": (
                [] if k_ar_params == 0 else constrain(np.arange(k_ar_params) / 10)
            ),
            "ma_params": (
                []
                if k_ma_params == 0
                else constrain((np.arange(k_ma_params) + 10) / 100)
            ),
            "seasonal_ar_params": (
                []
                if k_seasonal_ar_params == 0
                else constrain(np.arange(k_seasonal_ar_params) - 4)
            ),
            "seasonal_ma_params": (
                []
                if k_seasonal_ma_params == 0
                else constrain((np.arange(k_seasonal_ma_params) - 10) / 100)
            ),
            "sigma2": [] if concentrate_scale else 2.3424,
        }
    )

    # Test the spec created with order, seasonal_order
    spec = specification.SARIMAXSpecification(
        endog,
        exog=exog,
        order=(p, d, q),
        seasonal_order=(P, D, Q, s),
        enforce_stationarity=enforce_stationarity,
        enforce_invertibility=enforce_invertibility,
        concentrate_scale=concentrate_scale,
    )

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)

    # Test the spec created with ar_order, etc.
    spec = specification.SARIMAXSpecification(
        endog,
        exog=exog,
        ar_order=p,
        diff=d,
        ma_order=q,
        seasonal_ar_order=P,
        seasonal_diff=D,
        seasonal_ma_order=Q,
        seasonal_periods=s,
        enforce_stationarity=enforce_stationarity,
        enforce_invertibility=enforce_invertibility,
        concentrate_scale=concentrate_scale,
    )

    check_attributes(spec, *args, **kwargs)
    check_properties(spec, *args, **properties_kwargs)
    check_methods(spec, *args, **methods_kwargs)


def test_misc():
    # Check that no arguments results in all zero orders
    spec = specification.SARIMAXSpecification()
    assert_equal(spec.order, (0, 0, 0))
    assert_equal(spec.seasonal_order, (0, 0, 0, 0))

    # Check for repr
    spec = specification.SARIMAXSpecification(
        endog=pd.Series([0], name="y"),
        exog=pd.DataFrame([[0, 0]], columns=["x1", "x2"]),
        order=(1, 1, 2),
        seasonal_order=(2, 1, 0, 12),
        enforce_stationarity=False,
        enforce_invertibility=False,
        concentrate_scale=True,
    )
    desired = (
        "SARIMAXSpecification(endog=y, exog=['x1', 'x2'],"
        " order=(1, 1, 2), seasonal_order=(2, 1, 0, 12),"
        " enforce_stationarity=False, enforce_invertibility=False,"
        " concentrate_scale=True)"
    )
    assert_equal(repr(spec), desired)


def test_invalid():
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(1, 0, 0), ar_order=1)
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(
            seasonal_order=(1, 0, 0), seasonal_ar_order=1
        )
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(-1, 0, 0))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(1.5, 0, 0))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(0, -1, 0))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(0, 1.5, 0))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(0,))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(0, 1.5, 0, 4))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(-1, 0, 0, 4))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(1.5, 0, 0, 4))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(0, -1, 0, 4))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(0, 1.5, 0, 4))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(1, 0, 0, 0))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(1, 0, 0, -1))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(1, 0, 0, 1))
    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(seasonal_order=(1,))

    with pytest.raises(ValueError):
        specification.SARIMAXSpecification(order=(1, 0, 0), endog=np.zeros((10, 2)))

    spec = specification.SARIMAXSpecification(ar_order=1)
    with pytest.raises(ValueError):
        spec.join_params()
    with pytest.raises(ValueError):
        spec.join_params(ar_params=[0.2, 0.3])


@pytest.mark.parametrize(
    "order,seasonal_order,enforce_stationarity,"
    "enforce_invertibility,concentrate_scale,valid",
    [
        # Different orders
        (
            (0, 0, 0),
            (0, 0, 0, 0),
            None,
            None,
            None,
            [
                "yule_walker",
                "burg",
                "innovations",
                "hannan_rissanen",
                "innovations_mle",
                "statespace",
            ],
        ),
        (
            (1, 0, 0),
            (0, 0, 0, 0),
            None,
            None,
            None,
            ["yule_walker", "burg", "hannan_rissanen", "innovations_mle", "statespace"],
        ),
        (
            (0, 0, 1),
            (0, 0, 0, 0),
            None,
            None,
            None,
            ["innovations", "hannan_rissanen", "innovations_mle", "statespace"],
        ),
        (
            (1, 0, 1),
            (0, 0, 0, 0),
            None,
            None,
            None,
            ["hannan_rissanen", "innovations_mle", "statespace"],
        ),
        ((0, 0, 0), (1, 0, 0, 4), None, None, None, ["innovations_mle", "statespace"]),
        # Different options
        ((1, 0, 0), (0, 0, 0, 0), True, None, None, ["innovations_mle", "statespace"]),
        (
            (1, 0, 0),
            (0, 0, 0, 0),
            False,
            None,
            None,
            ["yule_walker", "burg", "hannan_rissanen", "statespace"],
        ),
        (
            (1, 0, 0),
            (0, 0, 0, 0),
            None,
            True,
            None,
            ["yule_walker", "burg", "hannan_rissanen", "innovations_mle", "statespace"],
        ),
        (
            (1, 0, 0),
            (0, 0, 0, 0),
            None,
            False,
            None,
            ["yule_walker", "burg", "hannan_rissanen", "innovations_mle", "statespace"],
        ),
        (
            (1, 0, 0),
            (0, 0, 0, 0),
            None,
            None,
            True,
            ["yule_walker", "burg", "hannan_rissanen", "statespace"],
        ),
    ],
)
def test_valid_estimators(
    order,
    seasonal_order,
    enforce_stationarity,
    enforce_invertibility,
    concentrate_scale,
    valid,
):
    # Basic specification
    spec = specification.SARIMAXSpecification(
        order=order,
        seasonal_order=seasonal_order,
        enforce_stationarity=enforce_stationarity,
        enforce_invertibility=enforce_invertibility,
        concentrate_scale=concentrate_scale,
    )

    estimators = {
        "yule_walker",
        "burg",
        "innovations",
        "hannan_rissanen",
        "innovations_mle",
        "statespace",
    }
    desired = set(valid)
    assert_equal(spec.valid_estimators, desired)
    for estimator in desired:
        assert_equal(spec.validate_estimator(estimator), None)
    for estimator in estimators.difference(desired):
        print(estimator, enforce_stationarity)
        with pytest.raises(ValueError):
            spec.validate_estimator(estimator)

    # Now try specification with missing values in endog
    spec = specification.SARIMAXSpecification(
        endog=[np.nan],
        order=order,
        seasonal_order=seasonal_order,
        enforce_stationarity=enforce_stationarity,
        enforce_invertibility=enforce_invertibility,
        concentrate_scale=concentrate_scale,
    )

    assert_equal(spec.valid_estimators, {"statespace"})
    assert_equal(spec.validate_estimator("statespace"), None)
    for estimator in estimators.difference(["statespace"]):
        with pytest.raises(ValueError):
            spec.validate_estimator(estimator)


def test_invalid_estimator():
    spec = specification.SARIMAXSpecification()
    with pytest.raises(ValueError):
        spec.validate_estimator("not_an_estimator")
