"""
Tests for fixing the values of some parameters and estimating others

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

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

from statsmodels import datasets
from statsmodels.tsa.statespace import (
    dynamic_factor,
    initialization,
    mlemodel,
    sarimax,
    structural,
    varmax,
)

macrodata = datasets.macrodata.load_pandas().data


def test_fix_params():
    # Just create a dummy model to test the basic `fix_params` mechanics
    mod = mlemodel.MLEModel([], 1)
    mod._param_names = ["a", "b", "c"]
    with mod.fix_params({"b": 1.0}):
        assert_(mod._has_fixed_params)
        assert_equal(mod._fixed_params, {"b": 1.0})
        assert_equal(mod._fixed_params_index, [1])
        assert_equal(mod._free_params_index, [0, 2])
    assert_(not mod._has_fixed_params)
    assert_equal(mod._fixed_params, {})
    assert_equal(mod._fixed_params_index, None)
    assert_equal(mod._free_params_index, None)


def test_nested_fix_params():
    mod = mlemodel.MLEModel([], 1)
    mod._param_names = ["a", "b", "c"]
    with mod.fix_params({"a": 2, "b": 0}):
        with mod.fix_params({"b": 1.0}):
            assert_(mod._has_fixed_params)
            assert_equal(mod._fixed_params, {"a": 2, "b": 1.0})
            assert_equal(mod._fixed_params_index, [0, 1])
            assert_equal(mod._free_params_index, [2])
    assert_(not mod._has_fixed_params)
    assert_equal(mod._fixed_params, {})
    assert_equal(mod._fixed_params_index, None)
    assert_equal(mod._free_params_index, None)


def test_results_append():
    endog = macrodata["infl"]
    endog1 = endog.iloc[:100]
    endog2 = endog.iloc[100:]

    mod_full = sarimax.SARIMAX(endog)
    with mod_full.fix_params({"ar.L1": 0.5}):
        res_full = mod_full.smooth([1.0], includes_fixed=False)
        # Start pretty close to optimum to speed up test
        start_params = [10.3]
        res_full_fit = mod_full.fit(start_params, disp=False)

    mod = sarimax.SARIMAX(endog1)
    with mod.fix_params({"ar.L1": 0.5}):
        res1 = mod.smooth([1.0], includes_fixed=False)

    # Append should work outside the context manager, since the results object
    # was created with fixed parameters
    res2 = res1.append(endog2)

    # Test to make sure refit works with fixed parameters
    res2_fit = res1.append(
        endog2,
        refit=True,
        fit_kwargs={"disp": False, "start_params": res_full_fit.params},
    )

    # Check non-refit
    assert_allclose(res2.params, res_full.params)
    assert_equal(res2._fixed_params, res_full._fixed_params)
    assert_allclose(res2.llf_obs, res_full.llf_obs)

    # Check refit results
    assert_allclose(res2_fit.params, res_full_fit.params)
    assert_equal(res2_fit._fixed_params, res_full_fit._fixed_params)
    assert_allclose(res2_fit.llf_obs, res_full_fit.llf_obs)


def test_results_extend():
    endog = macrodata["infl"]
    endog1 = endog.iloc[:100]
    endog2 = endog.iloc[100:]

    mod_full = sarimax.SARIMAX(endog)
    with mod_full.fix_params({"ar.L1": 0.5}):
        res_full = mod_full.smooth([1.0], includes_fixed=False)

    mod = sarimax.SARIMAX(endog1)
    with mod.fix_params({"ar.L1": 0.5}):
        res1 = mod.smooth([1.0], includes_fixed=False)

    # Append should work outside the context manager, since the results object
    # was created with fixed parameters
    res2 = res1.append(endog2)

    # Check results
    assert_allclose(res2.params, res_full.params)
    assert_equal(res2._fixed_params, res_full._fixed_params)
    assert_allclose(res2.llf_obs, res_full.llf_obs)


def test_results_apply():
    endog = macrodata["infl"]

    mod = sarimax.SARIMAX(endog)
    with mod.fix_params({"ar.L1": 0.5}):
        res = mod.smooth([1.0], includes_fixed=False)
        # Start pretty close to optimum to speed up test
        start_params = [10.3]
        res_fit = mod.fit(start_params, disp=False)

    res2 = res.apply(endog)

    # Test to make sure refit works with fixed parameters
    res2_fit = res.apply(
        endog, refit=True, fit_kwargs={"disp": False, "start_params": res_fit.params}
    )

    # Check non-refit
    assert_allclose(res2.params, res.params)
    assert_equal(res2._fixed_params, res._fixed_params)
    assert_allclose(res2.llf_obs, res.llf_obs)

    # Check refit results
    assert_allclose(res2_fit.params, res_fit.params)
    assert_equal(res2_fit._fixed_params, res_fit._fixed_params)
    assert_allclose(res2_fit.llf_obs, res_fit.llf_obs)


def test_mle_validate():
    mod = mlemodel.MLEModel([], 1)
    mod._param_names = ["a", "b", "c"]
    msg = 'Invalid parameter name passed: "d"'

    with pytest.raises(ValueError, match=msg):
        with mod.fix_params({"d": 1}):
            pass


def test_sarimax_validate():
    # Test for invalid uses of parameter fixing
    endog = macrodata["infl"]
    mod1 = sarimax.SARIMAX(endog, order=(2, 0, 0))

    # Try to fix invalid parameter
    with pytest.raises(ValueError):
        mod1.fit_constrained({"AR.L1": 0.5})

    # Cannot fix individual parameters that are part of a multivariate
    # transformation
    with pytest.raises(ValueError):
        with mod1.fix_params({"ar.L1": 0.5}):
            pass
    with pytest.raises(ValueError):
        mod1.fit_constrained({"ar.L1": 0.5})

    # But can fix the entire set of parameters that are part of a multivariate
    # transformation
    with mod1.fix_params({"ar.L1": 0.5, "ar.L2": 0.2}):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
        assert_equal(mod1._fixed_params_index, [0, 1])
        assert_equal(mod1._free_params_index, [2])

    res = mod1.fit_constrained(
        {"ar.L1": 0.5, "ar.L2": 0.2}, start_params=[7.0], disp=False
    )
    assert_(res._has_fixed_params)
    assert_equal(res._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
    assert_equal(res._fixed_params_index, [0, 1])
    assert_equal(res._free_params_index, [2])

    with mod1.fix_params({"ar.L1": 0.5, "ar.L2": 0.0}):
        # overwrite ar.L2 with nested fix_params
        with mod1.fix_params({"ar.L2": 0.2}):
            assert_(mod1._has_fixed_params)
            assert_equal(mod1._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
            assert_equal(mod1._fixed_params_index, [0, 1])
            assert_equal(mod1._free_params_index, [2])


def test_structural_validate():
    # Test for invalid uses of parameter fixing
    endog = macrodata["infl"]
    mod1 = structural.UnobservedComponents(endog, "rwalk", autoregressive=2)

    # Try to fix invalid parameter
    with pytest.raises(ValueError):
        mod1.fit_constrained({"AR.L1": 0.5})

    # Cannot fix parameters that are part of a multivariate transformation
    with pytest.raises(ValueError):
        with mod1.fix_params({"ar.L1": 0.5}):
            pass
    with pytest.raises(ValueError):
        mod1.fit_constrained({"ar.L1": 0.5})

    # But can fix the entire set of parameters that are part of a multivariate
    # transformation
    with mod1.fix_params({"ar.L1": 0.5, "ar.L2": 0.2}):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
        assert_equal(mod1._fixed_params_index, [2, 3])
        assert_equal(mod1._free_params_index, [0, 1])

    res = mod1.fit_constrained(
        {"ar.L1": 0.5, "ar.L2": 0.2}, start_params=[7.0], disp=False
    )
    assert_(res._has_fixed_params)
    assert_equal(res._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
    assert_equal(res._fixed_params_index, [2, 3])
    assert_equal(res._free_params_index, [0, 1])

    with mod1.fix_params({"ar.L1": 0.5, "ar.L2": 0.0}):
        # overwrite ar.L2 with nested fix_params
        with mod1.fix_params({"ar.L2": 0.2}):
            assert_(mod1._has_fixed_params)
            assert_equal(mod1._fixed_params, {"ar.L1": 0.5, "ar.L2": 0.2})
            assert_equal(mod1._fixed_params_index, [2, 3])
            assert_equal(mod1._free_params_index, [0, 1])


def test_dynamic_factor_validate():
    # Test for invalid uses of parameter fixing
    endog = np.log(macrodata[["cpi", "realgdp", "realinv"]]).diff().iloc[1:]
    endog = (endog - endog.mean()) / endog.std()

    # Basic model
    mod1 = dynamic_factor.DynamicFactor(
        endog, k_factors=1, factor_order=1, error_cov_type="diagonal"
    )

    # Check loading
    constraints = {"loading.f1.cpi": 0.5}
    with mod1.fix_params(constraints):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, constraints)
        assert_equal(mod1._fixed_params_index, [0])
        assert_equal(mod1._free_params_index, [1, 2, 3, 4, 5, 6])
    res1 = mod1.fit_constrained(constraints, disp=False)
    assert_(res1._has_fixed_params)
    assert_equal(res1._fixed_params, constraints)
    assert_equal(res1._fixed_params_index, [0])
    assert_equal(res1._free_params_index, [1, 2, 3, 4, 5, 6])

    # With k_factors=1 and factor_order=1, we can fix the factor AR coefficient
    # even with `enforce_stationarity=True`.
    # Fix factor AR coefficient
    with mod1.fix_params({"L1.f1.f1": 0.5}):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, {"L1.f1.f1": 0.5})
        assert_equal(mod1._fixed_params_index, [6])
        assert_equal(mod1._free_params_index, [0, 1, 2, 3, 4, 5])

    # With k_factors > 1 or factor_order > 1, we can only fix the entire set of
    # factor AR coefficients when `enforce_stationarity=True`.
    mod2 = dynamic_factor.DynamicFactor(
        endog, k_factors=1, factor_order=2, error_cov_type="diagonal"
    )
    with pytest.raises(ValueError):
        with mod2.fix_params({"L1.f1.f1": 0.5}):
            pass

    constraints = {"L1.f1.f1": 0.3, "L2.f1.f1": 0.1}
    with mod2.fix_params(constraints):
        assert_(mod2._has_fixed_params)
        assert_equal(mod2._fixed_params, constraints)
        assert_equal(mod2._fixed_params_index, [6, 7])
        assert_equal(mod2._free_params_index, [0, 1, 2, 3, 4, 5])

    res2 = mod2.fit_constrained(constraints, disp=False)
    assert_(res2._has_fixed_params)
    assert_equal(res2._fixed_params, constraints)
    assert_equal(res2._fixed_params_index, [6, 7])
    assert_equal(res2._free_params_index, [0, 1, 2, 3, 4, 5])

    with mod2.fix_params(constraints):
        # overwrite L1.f1.f1 with nested fix_params
        with mod2.fix_params({"L1.f1.f1": -0.3}):
            assert_(mod2._has_fixed_params)
            assert_equal(mod2._fixed_params, {"L1.f1.f1": -0.3, "L2.f1.f1": 0.1})
            assert_equal(mod2._fixed_params_index, [6, 7])
            assert_equal(mod2._free_params_index, [0, 1, 2, 3, 4, 5])

    # (same as previous, now k_factors=2)
    mod3 = dynamic_factor.DynamicFactor(
        endog, k_factors=2, factor_order=1, error_cov_type="diagonal"
    )
    with pytest.raises(ValueError):
        with mod3.fix_params({"L1.f1.f1": 0.3}):
            pass

    constraints = {
        "L1.f1.f1": 0.3, "L1.f2.f1": 0.1, "L1.f1.f2": -0.05, "L1.f2.f2": 0.1
    }
    with mod3.fix_params(constraints):
        assert_(mod3._has_fixed_params)
        assert_equal(mod3._fixed_params, constraints)
        assert_equal(mod3._fixed_params_index, [9, 10, 11, 12])
        assert_equal(mod3._free_params_index, [0, 1, 2, 3, 4, 5, 6, 7, 8])

    res3 = mod3.fit_constrained(constraints, disp=False)
    assert_(res3._has_fixed_params)
    assert_equal(res3._fixed_params, constraints)
    assert_equal(res3._fixed_params_index, [9, 10, 11, 12])
    assert_equal(res3._free_params_index, [0, 1, 2, 3, 4, 5, 6, 7, 8])

    with mod3.fix_params(constraints):
        # overwrite L1.f1.f1 and L1.f2.f2 with nested fix_params
        with mod3.fix_params({"L1.f1.f1": -0.3, "L1.f2.f2": -0.1}):
            assert_(mod3._has_fixed_params)
            assert_equal(
                mod3._fixed_params,
                {
                    "L1.f1.f1": -0.3,
                    "L1.f2.f1": 0.1,
                    "L1.f1.f2": -0.05,
                    "L1.f2.f2": -0.1,
                },
            )
            assert_equal(mod3._fixed_params_index, [9, 10, 11, 12])
            assert_equal(mod3._free_params_index, [0, 1, 2, 3, 4, 5, 6, 7, 8])

    # Now, with enforce_stationarity=False, we can fix any of the factor AR
    # coefficients
    mod4 = dynamic_factor.DynamicFactor(
        endog,
        k_factors=1,
        factor_order=2,
        error_cov_type="diagonal",
        enforce_stationarity=False,
    )
    with mod4.fix_params({"L1.f1.f1": 0.6}):
        assert_(mod4._has_fixed_params)
        assert_equal(mod4._fixed_params, {"L1.f1.f1": 0.6})
        assert_equal(mod4._fixed_params_index, [6])
        assert_equal(mod4._free_params_index, [0, 1, 2, 3, 4, 5, 7])

    mod5 = dynamic_factor.DynamicFactor(
        endog,
        k_factors=2,
        factor_order=1,
        error_cov_type="diagonal",
        enforce_stationarity=False,
    )
    with mod5.fix_params({"L1.f1.f1": 0.6}):
        assert_(mod5._has_fixed_params)
        assert_equal(mod5._fixed_params, {"L1.f1.f1": 0.6})
        assert_equal(mod5._fixed_params_index, [9])
        assert_equal(mod5._free_params_index, [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12])

    # Check error variance
    # (mod1 has error_cov_type='diagonal', so we can fix any that we like)
    constraints = {"sigma2.cpi": 0.9, "sigma2.realinv": 3}
    with mod1.fix_params(constraints):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, constraints)
        assert_equal(mod1._fixed_params_index, [3, 5])
        assert_equal(mod1._free_params_index, [0, 1, 2, 4, 6])
    res1 = mod1.fit_constrained(constraints, disp=False)
    assert_(res1._has_fixed_params)
    assert_equal(res1._fixed_params, constraints)
    assert_equal(res1._fixed_params_index, [3, 5])
    assert_equal(res1._free_params_index, [0, 1, 2, 4, 6])

    # Check unstructured error variance
    # (also reduce k_endog and fix some other parameters to make MLE go faster)
    mod6 = dynamic_factor.DynamicFactor(
        endog[["cpi", "realgdp"]],
        k_factors=1,
        factor_order=1,
        error_cov_type="unstructured",
    )
    constraints = {
        "loading.f1.cpi": 1.0,
        "loading.f1.realgdp": 1.0,
        "cov.chol[1,1]": 0.5,
        "cov.chol[2,1]": 0.1,
    }
    with mod6.fix_params(constraints):
        assert_(mod6._has_fixed_params)
        assert_equal(mod6._fixed_params, constraints)
        assert_equal(mod6._fixed_params_index, [0, 1, 2, 3])
        assert_equal(mod6._free_params_index, [4, 5])
    res6 = mod6.fit_constrained(constraints, disp=False)
    assert_(res6._has_fixed_params)
    assert_equal(res6._fixed_params, constraints)
    assert_equal(res6._fixed_params_index, [0, 1, 2, 3])
    assert_equal(res6._free_params_index, [4, 5])


def test_varmax_validate():
    # Test for invalid uses of parameter fixing
    endog = np.log(macrodata[["cpi", "realgdp"]]).diff().iloc[1:]
    exog = np.log(macrodata[["realinv"]]).diff().iloc[1:]

    # Basic model, VAR(1) + exog + measurement error
    mod1 = varmax.VARMAX(endog, order=(1, 0), exog=exog, measurement_error=True)

    # Check intercept, exog
    constraints = {
        "intercept.cpi": 0.5,
        "intercept.realgdp": 1.1,
        "beta.realinv.cpi": 0.2,
        "beta.realinv.realgdp": 0.1,
        "sqrt.var.cpi": 1.2,
        "sqrt.cov.cpi.realgdp": -0.1,
        "sqrt.var.realgdp": 2.3,
        "measurement_variance.cpi": 0.4,
        "measurement_variance.realgdp": 0.4,
    }
    with mod1.fix_params(constraints):
        assert_(mod1._has_fixed_params)
        assert_equal(mod1._fixed_params, constraints)
        assert_equal(mod1._fixed_params_index, [0, 1, 6, 7, 8, 9, 10, 11, 12])
        assert_equal(mod1._free_params_index, [2, 3, 4, 5])
    res1 = mod1.fit_constrained(constraints, disp=False)
    assert_(res1._has_fixed_params)
    assert_equal(res1._fixed_params, constraints)
    assert_equal(res1._fixed_params_index, [0, 1, 6, 7, 8, 9, 10, 11, 12])
    assert_equal(res1._free_params_index, [2, 3, 4, 5])

    # With k_endog=1 we can fix the factor AR coefficient
    # even with `enforce_stationarity=True`.
    # Fix factor AR coefficient
    mod2 = varmax.VARMAX(
        endog[["cpi"]], order=(1, 0), exog=exog, measurement_error=True
    )
    constraints = {"L1.cpi.cpi": 0.5}
    with mod2.fix_params(constraints):
        assert_(mod2._has_fixed_params)
        assert_equal(mod2._fixed_params, constraints)
        assert_equal(mod2._fixed_params_index, [1])
        assert_equal(mod2._free_params_index, [0, 2, 3, 4])

    # With k_ar > 1, we can only fix the entire set of AR coefficients when
    # `enforce_stationarity=True`.
    mod3 = varmax.VARMAX(endog[["cpi"]], order=(2, 0))
    with pytest.raises(ValueError):
        with mod3.fix_params({"L1.cpi.cpi": 0.5}):
            pass

    constraints = {"L1.cpi.cpi": 0.3, "L2.cpi.cpi": 0.1}
    with mod3.fix_params(constraints):
        assert_(mod3._has_fixed_params)
        assert_equal(mod3._fixed_params, constraints)
        assert_equal(mod3._fixed_params_index, [1, 2])
        assert_equal(mod3._free_params_index, [0, 3])

    res3 = mod3.fit_constrained(constraints, start_params=[0, 1.0], disp=False)
    assert_(res3._has_fixed_params)
    assert_equal(res3._fixed_params, constraints)
    assert_equal(res3._fixed_params_index, [1, 2])
    assert_equal(res3._free_params_index, [0, 3])

    with mod3.fix_params(constraints):
        # overwrite L1.cpi.cpi with nested fix_params
        with mod3.fix_params({"L1.cpi.cpi": -0.3}):
            assert_(mod3._has_fixed_params)
            assert_equal(mod3._fixed_params, {"L1.cpi.cpi": -0.3, "L2.cpi.cpi": 0.1})
            assert_equal(mod3._fixed_params_index, [1, 2])
            assert_equal(mod3._free_params_index, [0, 3])

    # With k_endog > 1, we can only fix the entire set of AR coefficients when
    # `enforce_stationarity=True`.
    mod4 = varmax.VARMAX(endog, order=(1, 0))
    with pytest.raises(ValueError):
        with mod4.fix_params({"L1.cpi.cpi": 0.3}):
            pass
    constraints = {
        "L1.cpi.cpi": 0.3,
        "L1.realgdp.cpi": 0.1,
        "L1.cpi.realgdp": -0.05,
        "L1.realgdp.realgdp": 0.1,
    }
    with mod4.fix_params(constraints):
        assert_(mod4._has_fixed_params)
        assert_equal(mod4._fixed_params, constraints)
        assert_equal(mod4._fixed_params_index, [2, 3, 4, 5])
        assert_equal(mod4._free_params_index, [0, 1, 6, 7, 8])
    res4 = mod4.fit_constrained(constraints, disp=False)
    assert_(res4._has_fixed_params)
    assert_equal(res4._fixed_params, constraints)
    assert_equal(res4._fixed_params_index, [2, 3, 4, 5])
    assert_equal(res4._free_params_index, [0, 1, 6, 7, 8])

    # Now, with enforce_stationarity=False, we can fix any of the AR
    # coefficients
    mod5 = varmax.VARMAX(endog[["cpi"]], order=(1, 0), enforce_stationarity=False)
    with mod5.fix_params({"L1.cpi.cpi": 0.6}):
        assert_(mod5._has_fixed_params)
        assert_equal(mod5._fixed_params, {"L1.cpi.cpi": 0.6})
        assert_equal(mod5._fixed_params_index, [1])
        assert_equal(mod5._free_params_index, [0, 2])

    # Now, with enforce_stationarity=False, we can fix any of the AR
    # coefficients
    mod6 = varmax.VARMAX(endog, order=(1, 0), enforce_stationarity=False)
    with mod6.fix_params({"L1.cpi.cpi": 0.6}):
        assert_(mod6._has_fixed_params)
        assert_equal(mod6._fixed_params, {"L1.cpi.cpi": 0.6})
        assert_equal(mod6._fixed_params_index, [2])
        assert_equal(mod6._free_params_index, [0, 1, 3, 4, 5, 6, 7, 8])


def check_results(res1, res2, check_lutkepohl=False, check_params=True):
    # Check other results
    assert_allclose(res2.nobs, res1.nobs)
    assert_allclose(res2.nobs_diffuse, res1.nobs_diffuse)
    assert_allclose(res2.nobs_effective, res1.nobs_effective)
    assert_allclose(res2.k_diffuse_states, res1.k_diffuse_states)

    assert_allclose(res2.df_model, res1.df_model)
    assert_allclose(res2.df_resid, res1.df_resid)

    assert_allclose(res2.llf, res1.llf)
    assert_allclose(res2.aic, res1.aic)
    assert_allclose(res2.bic, res1.bic)
    assert_allclose(res2.hqic, res1.hqic)

    if check_lutkepohl:
        assert_allclose(
            res2.info_criteria("aic", "lutkepohl"),
            res1.info_criteria("aic", "lutkepohl"),
        )
        assert_allclose(
            res2.info_criteria("bic", "lutkepohl"),
            res1.info_criteria("bic", "lutkepohl"),
        )
        assert_allclose(
            res2.info_criteria("hqic", "lutkepohl"),
            res1.info_criteria("hqic", "lutkepohl"),
        )

    assert_allclose(res2.llf_obs, res1.llf_obs)
    assert_allclose(res2.fittedvalues, res1.fittedvalues)
    assert_allclose(res2.fittedvalues, res1.fittedvalues)

    if check_params:
        # Check parameter-related values
        mask_free = res2._free_params_index
        mask_fixed = res2._fixed_params_index
        pvalues = np.asarray(res2.pvalues)
        assert_allclose(pvalues[mask_free], res1.pvalues)
        assert_allclose(pvalues[mask_fixed], np.nan)

        bse = np.asarray(res2.bse)
        assert_allclose(bse[mask_free], res1.bse)
        assert_allclose(bse[mask_fixed], np.nan)

        zvalues = np.asarray(res2.zvalues)
        assert_allclose(zvalues[mask_free], res1.zvalues)
        assert_allclose(zvalues[mask_fixed], np.nan)

        # Check parameter covariance matrix
        mask_free = np.ix_(res2._free_params_index, res2._free_params_index)
        mask_fixed = np.ix_(res2._fixed_params_index, res2._fixed_params_index)
        assert_allclose(
            res2.cov_params_default.values[mask_free], res1.cov_params_default
        )
        assert_allclose(res2.cov_params_default.values[mask_fixed], np.nan)

        assert_allclose(
            res2.cov_params_approx.values[mask_free], res1.cov_params_approx
        )
        assert_allclose(res2.cov_params_approx.values[mask_fixed], np.nan)

        assert_allclose(res2.cov_params_oim.values[mask_free], res1.cov_params_oim)
        assert_allclose(res2.cov_params_oim.values[mask_fixed], np.nan)

        assert_allclose(res2.cov_params_opg.values[mask_free], res1.cov_params_opg)
        assert_allclose(res2.cov_params_opg.values[mask_fixed], np.nan)

        assert_allclose(
            res2.cov_params_robust.values[mask_free], res1.cov_params_robust
        )
        assert_allclose(res2.cov_params_robust.values[mask_fixed], np.nan)

        assert_allclose(
            res2.cov_params_robust_oim.values[mask_free], res1.cov_params_robust_oim
        )
        assert_allclose(res2.cov_params_robust_oim.values[mask_fixed], np.nan)

        assert_allclose(
            res2.cov_params_robust_approx.values[mask_free],
            res1.cov_params_robust_approx,
        )
        assert_allclose(res2.cov_params_robust_approx.values[mask_fixed], np.nan)

    # Check residual hypothesis tests
    assert_allclose(
        res2.test_normality("jarquebera"), res1.test_normality("jarquebera")
    )

    assert_allclose(
        res2.test_heteroskedasticity("breakvar"),
        res1.test_heteroskedasticity("breakvar"),
    )

    actual = res2.test_serial_correlation("ljungbox")
    desired = res1.test_serial_correlation("ljungbox")
    assert_allclose(actual, desired)


def test_sarimax_nonconsecutive():
    # SARIMAX allows using non-consecutive lag orders, which implicitly fix
    # AR coefficients to zeros, so we can test explicitly fixed AR coefficients
    # against this
    endog = macrodata["infl"]

    # y_t = \phi_1 y_{t-1} + \phi_4 y_{t-4} + \varepsilon_t
    # Note: because the transformation will not respect the parameter
    # constraints, we will need to set enforce_stationarity=False; set it to
    # False here too so that they both are the same model
    mod1 = sarimax.SARIMAX(
        endog, order=([1, 0, 0, 1], 0, 0), enforce_stationarity=False
    )
    mod2 = sarimax.SARIMAX(endog, order=(4, 0, 0), enforce_stationarity=False)

    # Start pretty close to optimum to speed up test
    start_params = [0.6, 0.2, 6.4]
    res1 = mod1.fit(start_params, disp=False)
    res2 = mod2.fit_constrained(
        {"ar.L2": 0, "ar.L3": 0}, res1.params, includes_fixed=False, disp=False
    )

    # Check that the right parameters were fixed
    assert_equal(res1.fixed_params, [])
    assert_equal(res2.fixed_params, ["ar.L2", "ar.L3"])

    # Check that MLE finds the same parameters in either case
    params = np.asarray(res1.params)
    desired = np.r_[params[0], 0, 0, params[1:]]
    assert_allclose(res2.params, desired)

    # Now smooth at the actual parameters (to allow high precision testing
    # below, even if there are small differences between MLE fitted parameters)
    with mod2.fix_params({"ar.L2": 0, "ar.L3": 0}):
        res2 = mod2.smooth(res1.params)

    check_results(res1, res2, check_lutkepohl=True)

    # Check that results methods work within the context
    with mod2.fix_params({"ar.L2": 0, "ar.L3": 0}):
        res3 = mod2.filter(res2.params, includes_fixed=True)
        check_results(res1, res3, check_lutkepohl=True)


def test_structural():
    # Many of the forms of UnobservedComponents are just special cases of the
    # most general model with parameter restrictions
    endog = macrodata["infl"]

    # Local linear trend is a pretty general model, so we can test fixing
    # parameters against other more specific models

    # Local level is local linear trend with sigma2.trend = 0
    mod1 = structural.UnobservedComponents(endog, "llevel")
    mod2 = structural.UnobservedComponents(endog, "lltrend")

    # Note: have to reinitialize, so the estimate of the trend term stays at
    # zero.
    init = initialization.Initialization(mod2.k_states)
    init[0] = "approximate_diffuse"
    init.set(1, "known", constant=[0])
    mod2.ssm.initialization = init
    mod2.ssm.loglikelihood_burn = 1

    constraints = {"sigma2.trend": 0}

    # Start pretty close to optimum to speed up test
    start_params = [3.37, 0.74]
    res1 = mod1.fit(start_params, disp=False)
    res2 = mod2.fit_constrained(
        constraints, start_params=res1.params, includes_fixed=False, disp=False
    )

    # Check that the right parameters were fixed
    assert_equal(res1.fixed_params, [])
    assert_equal(res2.fixed_params, ["sigma2.trend"])

    # Check that MLE finds the same parameters in either case
    desired = np.r_[res1.params, 0]
    assert_allclose(res2.params, desired)

    # Now smooth at the actual parameters (to allow high precision testing
    # below, even if there are small differences between MLE fitted parameters)
    with mod2.fix_params(constraints):
        res2 = mod2.smooth(res1.params)

    check_results(res1, res2)


def test_dynamic_factor_diag_error_cov():
    # Can test fixing the off-diagonal error covariance parameters to zeros
    # with `error_cov_type='unstructured'` against the case
    # `error_cov_type='diagonal'`.
    endog = np.log(macrodata[["cpi", "realgdp"]]).diff().iloc[1:]
    endog = (endog - endog.mean()) / endog.std()

    # Basic model
    mod1 = dynamic_factor.DynamicFactor(
        endog, k_factors=1, factor_order=1, error_cov_type="diagonal"
    )
    mod2 = dynamic_factor.DynamicFactor(
        endog, k_factors=1, factor_order=1, error_cov_type="unstructured"
    )

    constraints = {"cov.chol[2,1]": 0}

    # Start pretty close to optimum to speed up test
    start_params = [-4.5e-06, -1.0e-05, 9.9e-01, 9.9e-01, -1.4e-01]
    res1 = mod1.fit(start_params=start_params, disp=False)
    res2 = mod2.fit_constrained(
        constraints, start_params=res1.params, includes_fixed=False, disp=False
    )

    # Check that the right parameters were fixed
    assert_equal(res1.fixed_params, [])
    assert_equal(res2.fixed_params, ["cov.chol[2,1]"])

    # Check that MLE finds the same parameters in either case
    # (need to account for the fact that diagonal params are variances but
    # unstructured params are standard deviations)
    param_vals = np.asarray(res1.params)
    params = np.r_[param_vals[:2], param_vals[2:4] ** 0.5, param_vals[4]]
    desired = np.r_[params[:3], 0, params[3:]]
    assert_allclose(res2.params, desired, atol=1e-5)

    # Now smooth at the actual parameters (to allow high precision testing
    # below, even if there are small differences between MLE fitted parameters)
    with mod2.fix_params(constraints):
        res2 = mod2.smooth(params)

    # Can't check some parameters-related values because of the different
    # parameterization (i.e. cov_params, bse, pvalues, etc. won't match).
    check_results(res1, res2, check_params=False)


def test_score_shape():
    # Test that the `score()` output for fixed params has the same shape as the
    # input vector
    endog = macrodata["infl"]

    mod = sarimax.SARIMAX(endog, order=(1, 0, 0))

    with mod.fix_params({"ar.L1": 0.5}):
        score = mod.score([1.0])

    assert_equal(score.shape, (1,))
