"""
Basic tests for the BM model

These test standard that the state space model is set up as desired, that
expected exceptions are raised, etc.
"""
from statsmodels.compat.pandas import (
    MONTH_END,
    QUARTER_END,
    assert_frame_equal,
    assert_series_equal,
)

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

from statsmodels.regression.linear_model import OLS
from statsmodels.tools import add_constant
from statsmodels.tsa.statespace import (
    dynamic_factor,
    dynamic_factor_mq,
    sarimax,
)
from statsmodels.tsa.statespace.tests import test_dynamic_factor_mq_monte_carlo

SKIP_MONTE_CARLO_TESTS = True


def test_default():
    # Includes 2 monthly and 2 quarterly series
    # Default is factors=1, factor_orders=1, idiosyncratic_ar1=True

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=1, factor_orders=1,
        idiosyncratic_ar1=True)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 5 + 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 1 + 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "L1.0", "L2.0", "L3.0", "L4.0"] +
               ["eps_M.M0", "eps_M.M1", "eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.0->M1", "loading.0->Q0",
        "loading.0->Q1",
        "L1.0->0", "fb(0).cov.chol[1,1]",
        "L1.eps_M.M0", "L1.eps_M.M1",
        "L1.eps_Q.Q0", "L1.eps_Q.Q1",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", :2, 5:7], np.eye(2))
    assert_allclose(mod["design", 2:, 7:9], np.eye(2))
    assert_allclose(mod["design", 2:, 9:11], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 11:13], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 13:15], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 15:17], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 20)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 1:5, :4], np.eye(4))
    assert_allclose(mod["transition", 9:17, 7:15], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 12)

    assert_allclose(mod["selection", 0, 0], np.eye(1))
    assert_allclose(mod["selection", 5:7, 1:3], np.eye(2))
    assert_allclose(mod["selection", 7:9, 3:5], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 5)

    assert_allclose(mod["state_cov"], 0)

    # Test parameter entry
    mod.update(np.arange(mod.k_params) + 2)

    # -> obs_intercept
    assert_allclose(mod["obs_intercept"], 0)

    # -> design
    desired = np.array([
        [2., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [3., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [4., 8., 12, 8., 4., 0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1., 0.],
        [5., 10, 15, 10, 5., 0., 0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1.]])
    assert_allclose(mod["design"], desired)

    # -> obs_cov
    assert_allclose(mod["obs_cov"], 0)

    # -> state_intercept
    assert_allclose(mod["state_intercept"], 0)

    # -> transition
    desired = np.array([
        [6., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 8., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 9., 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 11., 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1., 0, 0]])
    assert_allclose(mod["transition"], desired)

    # -> selection
    assert_allclose(np.sum(mod["selection"]), 5)

    # -> state_cov
    desired = np.array([[49., 0., 0., 0., 0.],
                        [0., 12., 0., 0., 0.],
                        [0., 0., 13., 0., 0.],
                        [0., 0., 0., 14., 0.],
                        [0., 0., 0., 0., 15.]])
    assert_allclose(mod["state_cov"], desired)


def test_k_factors_gt1():
    # Includes 2 monthly and 2 quarterly series
    # This case: k_factors=2, factor_orders=1, idiosyncratic_ar1=True

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=2, factor_orders={("0", "1"): 1},
        idiosyncratic_ar1=True)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 5 * 2 + 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 2 + 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "1", "L1.0", "L1.1", "L2.0", "L2.1",
                "L3.0", "L3.1", "L4.0", "L4.1"] +
               ["eps_M.M0", "eps_M.M1", "eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.1->M0", "loading.0->M1",
        "loading.1->M1",
        "loading.0->Q0", "loading.1->Q0", "loading.0->Q1",
        "loading.1->Q1",
        "L1.0->0", "L1.1->0", "L1.0->1", "L1.1->1",
        "fb(0).cov.chol[1,1]", "fb(0).cov.chol[2,1]", "fb(0).cov.chol[2,2]",
        "L1.eps_M.M0", "L1.eps_M.M1",
        "L1.eps_Q.Q0", "L1.eps_Q.Q1",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", :2, 10:12], np.eye(2))
    assert_allclose(mod["design", 2:, 12:14], np.eye(2))
    assert_allclose(mod["design", 2:, 14:16], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 16:18], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 18:20], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 20:22], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 20)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 2:10, :8], np.eye(8))
    assert_allclose(mod["transition", 14:22, 12:20], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 16)

    assert_allclose(mod["selection", :2, :2], np.eye(2))
    assert_allclose(mod["selection", 10:12, 2:4], np.eye(2))
    assert_allclose(mod["selection", 12:14, 4:6], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 6)

    assert_allclose(mod["state_cov"], 0)

    # Test parameter entry
    mod.update(np.arange(mod.k_params) + 2)

    # -> obs_intercept
    assert_allclose(mod["obs_intercept"], 0)

    # -> design
    desired = np.array([
        [2., 3., 0., 0., 0., 0., 0., 0., 0., 0.],
        [4., 5., 0., 0., 0., 0., 0., 0., 0., 0.],
        [6., 7., 12, 14, 18, 21, 12, 14, 6., 7.],
        [8., 9., 16, 18, 24, 27, 16, 18, 8., 9.]])
    assert_allclose(mod["design", :, :10], desired)
    desired = np.array([
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1., 0.],
        [0., 0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1.]])
    assert_allclose(mod["design", :, 10:], desired)

    # -> obs_cov
    assert_allclose(mod["obs_cov"], 0)

    # -> state_intercept
    assert_allclose(mod["state_intercept"], 0)

    # -> transition
    desired = np.array([
        [10, 11, 0., 0., 0., 0., 0., 0., 0., 0.],
        [12, 13, 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", :10, :10], desired)
    assert_allclose(mod["transition", :10, 10:], 0)
    assert_allclose(mod["transition", 10:, :10], 0)
    desired = np.array([
        [17, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 18, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 19, 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 20, 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", 10:, 10:], desired)

    # -> selection
    assert_allclose(np.sum(mod["selection"]), 6)

    # -> state_cov
    L = np.array([[14., 0],
                  [15., 16.]])
    desired = np.array([[0., 0., 0., 0., 0., 0.],
                        [0., 0., 0., 0., 0., 0.],
                        [0., 0., 21, 0., 0., 0.],
                        [0., 0., 0., 22, 0., 0.],
                        [0., 0., 0., 0., 23, 0.],
                        [0., 0., 0., 0., 0., 24]])
    desired[:2, :2] = np.dot(L, L.T)
    assert_allclose(mod["state_cov"], desired)


def test_factor_order_gt1():
    # Includes 2 monthly and 2 quarterly series
    # This case: factors=1, factor_orders=6, idiosyncratic_ar1=True

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=1, factor_orders=6,
        idiosyncratic_ar1=True)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 6 + 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 1 + 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "L1.0", "L2.0", "L3.0", "L4.0",
                "L5.0"] +
               ["eps_M.M0", "eps_M.M1", "eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.0->M1", "loading.0->Q0",
        "loading.0->Q1",
        "L1.0->0", "L2.0->0", "L3.0->0", "L4.0->0",
        "L5.0->0", "L6.0->0",
        "fb(0).cov.chol[1,1]",
        "L1.eps_M.M0", "L1.eps_M.M1",
        "L1.eps_Q.Q0", "L1.eps_Q.Q1",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", :2, 6:8], np.eye(2))
    assert_allclose(mod["design", 2:, 8:10], np.eye(2))
    assert_allclose(mod["design", 2:, 10:12], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 12:14], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 14:16], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 16:18], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 20)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 1:6, :5], np.eye(5))
    assert_allclose(mod["transition", 10:18, 8:16], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 13)

    assert_allclose(mod["selection", 0, 0], np.eye(1))
    assert_allclose(mod["selection", 6:8, 1:3], np.eye(2))
    assert_allclose(mod["selection", 8:10, 3:5], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 5)

    assert_allclose(mod["state_cov"], 0)

    # Test parameter entry
    mod.update(np.arange(mod.k_params) + 2)

    # -> obs_intercept
    assert_allclose(mod["obs_intercept"], 0)

    # -> design
    desired = np.array([
        [2., 0., 0., 0., 0., 0.],
        [3., 0., 0., 0., 0., 0.],
        [4., 8., 12, 8., 4., 0.],
        [5., 10, 15, 10, 5., 0.]])
    assert_allclose(mod["design", :, :6], desired)
    desired = np.array([
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1., 0.],
        [0., 0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1.]])
    assert_allclose(mod["design", :, 6:], desired)

    # -> obs_cov
    assert_allclose(mod["obs_cov"], 0)

    # -> state_intercept
    assert_allclose(mod["state_intercept"], 0)

    # -> transition
    desired = np.array([
        [6., 7., 8., 9., 10, 11, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 13, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 14, 0., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 15, 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 16, 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0, 0],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0, 0]
    ])
    assert_allclose(mod["transition"], desired)

    # -> selection
    assert_allclose(np.sum(mod["selection"]), 5)

    # -> state_cov
    desired = np.array([[144, 0., 0., 0., 0.],
                        [0., 17., 0., 0., 0.],
                        [0., 0., 18., 0., 0.],
                        [0., 0., 0., 19., 0.],
                        [0., 0., 0., 0., 20.]])
    assert_allclose(mod["state_cov"], desired)


def test_factor_order_12():
    # Includes 2 monthly and 2 quarterly series
    # This case: factors= 2, factor_orders=(1,2), idiosyncratic_ar1=True

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=2, factor_orders={"0": 2},
        idiosyncratic_ar1=True)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 5 * 2 + 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 2 + 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "L1.0", "L2.0", "L3.0", "L4.0", "1",
                "L1.1", "L2.1", "L3.1", "L4.1"] +
               ["eps_M.M0", "eps_M.M1", "eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.1->M0", "loading.0->M1", "loading.1->M1",
        "loading.0->Q0", "loading.1->Q0", "loading.0->Q1", "loading.1->Q1",
        "L1.0->0", "L2.0->0", "L1.1->1",
        "fb(0).cov.chol[1,1]", "fb(1).cov.chol[1,1]",
        "L1.eps_M.M0", "L1.eps_M.M1",
        "L1.eps_Q.Q0", "L1.eps_Q.Q1",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", :2, 10:12], np.eye(2))
    assert_allclose(mod["design", 2:, 12:14], np.eye(2))
    assert_allclose(mod["design", 2:, 14:16], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 16:18], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 18:20], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 20:22], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 20)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 1:5, :4], np.eye(4))
    assert_allclose(mod["transition", 14:22, 12:20], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 16)

    assert_allclose(mod["selection", :2, :2], np.array([[1, 0], [0, 0]]))
    assert_allclose(mod["selection", 4:6, :2], np.array([[0, 0], [0, 1]]))
    assert_allclose(mod["selection", 10:12, 2:4], np.eye(2))
    assert_allclose(mod["selection", 12:14, 4:6], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 6)

    assert_allclose(mod["state_cov"], 0)


def test_k_factors_gt1_factor_order_gt1():
    # Includes 2 monthly and 2 quarterly series
    # This case: kactors=2, factor_orders=6, idiosyncratic_ar1=True

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=2, factor_orders={("0", "1"): 6},
        idiosyncratic_ar1=True)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 6 * 2 + 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 2 + 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "1", "L1.0", "L1.1", "L2.0", "L2.1",
                "L3.0", "L3.1", "L4.0", "L4.1", "L5.0",
                "L5.1"] +
               ["eps_M.M0", "eps_M.M1", "eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.1->M0", "loading.0->M1",
        "loading.1->M1",
        "loading.0->Q0", "loading.1->Q0", "loading.0->Q1",
        "loading.1->Q1",
        "L1.0->0", "L1.1->0", "L2.0->0", "L2.1->0",
        "L3.0->0", "L3.1->0", "L4.0->0", "L4.1->0",
        "L5.0->0", "L5.1->0", "L6.0->0", "L6.1->0",
        "L1.0->1", "L1.1->1", "L2.0->1", "L2.1->1",
        "L3.0->1", "L3.1->1", "L4.0->1", "L4.1->1",
        "L5.0->1", "L5.1->1", "L6.0->1", "L6.1->1",
        "fb(0).cov.chol[1,1]", "fb(0).cov.chol[2,1]", "fb(0).cov.chol[2,2]",
        "L1.eps_M.M0", "L1.eps_M.M1",
        "L1.eps_Q.Q0", "L1.eps_Q.Q1",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", :2, 12:14], np.eye(2))
    assert_allclose(mod["design", 2:, 14:16], np.eye(2))
    assert_allclose(mod["design", 2:, 16:18], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 18:20], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 20:22], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 22:24], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 20)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 2:12, :10], np.eye(10))
    assert_allclose(mod["transition", 16:24, 14:22], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 18)

    assert_allclose(mod["selection", :2, :2], np.eye(2))
    assert_allclose(mod["selection", 12:14, 2:4], np.eye(2))
    assert_allclose(mod["selection", 14:16, 4:6], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 6)

    assert_allclose(mod["state_cov"], 0)

    # Test parameter entry
    mod.update(np.arange(mod.k_params) + 2)

    # -> obs_intercept
    assert_allclose(mod["obs_intercept"], 0)

    # -> design
    desired = np.array([
        [2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [4., 5., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [6., 7., 12, 14, 18, 21, 12, 14, 6., 7., 0., 0.],
        [8., 9., 16, 18, 24, 27, 16, 18, 8., 9., 0., 0.]])
    assert_allclose(mod["design", :, :12], desired)
    desired = np.array([
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1., 0.],
        [0., 0., 0., 1., 0., 2., 0., 3., 0., 2., 0., 1.]])
    assert_allclose(mod["design", :, 12:], desired)

    # -> obs_cov
    assert_allclose(mod["obs_cov"], 0)

    # -> state_intercept
    assert_allclose(mod["state_intercept"], 0)

    # -> transition
    desired = np.array([
        [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
        [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", :12, :12], desired)
    assert_allclose(mod["transition", :12, 12:], 0)
    assert_allclose(mod["transition", 12:, :12], 0)
    desired = np.array([
        [37, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 38, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 39, 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 40, 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", 12:, 12:], desired)

    # -> selection
    assert_allclose(np.sum(mod["selection"]), 6)

    # -> state_cov
    L = np.array([[34., 0],
                  [35., 36.]])
    desired = np.array([[0., 0., 0., 0., 0., 0.],
                        [0., 0., 0., 0., 0., 0.],
                        [0., 0., 41, 0., 0., 0.],
                        [0., 0., 0., 42, 0., 0.],
                        [0., 0., 0., 0., 43, 0.],
                        [0., 0., 0., 0., 0., 44]])
    desired[:2, :2] = np.dot(L, L.T)
    assert_allclose(mod["state_cov"], desired)


def test_k_factors_gt1_factor_order_gt1_no_idiosyncratic_ar1():
    # Includes 2 monthly and 2 quarterly series
    # This case: factors=2, factor_orders=6, idiosyncratic_ar1=False

    # Create the datasets
    index_M = pd.period_range(start="2000", periods=12, freq="M")
    index_Q = pd.period_range(start="2000", periods=4, freq="Q")

    dta_M = pd.DataFrame(np.zeros((12, 2)), index=index_M,
                         columns=["M0", "M1"])
    dta_Q = pd.DataFrame(np.zeros((4, 2)), index=index_Q, columns=["Q0", "Q1"])
    # Add some noise so the variables aren't constants
    dta_M.iloc[0] = 1.
    dta_Q.iloc[1] = 1.

    # Create the model instance
    mod = dynamic_factor_mq.DynamicFactorMQ(
        dta_M, endog_quarterly=dta_Q, factors=2, factor_orders={("0", "1"): 6},
        idiosyncratic_ar1=False)

    # Test dimensions
    assert_equal(mod.k_endog, 2 + 2)
    assert_equal(mod.k_states, 6 * 2 + 2 * 5)
    assert_equal(mod.ssm.k_posdef, 2 + 2)

    # Test names
    assert_equal(mod.endog_names, ["M0", "M1", "Q0", "Q1"])
    desired = (["0", "1", "L1.0", "L1.1", "L2.0", "L2.1",
                "L3.0", "L3.1", "L4.0", "L4.1", "L5.0",
                "L5.1"] +
               ["eps_Q.Q0", "eps_Q.Q1"] +
               ["L1.eps_Q.Q0", "L1.eps_Q.Q1"] +
               ["L2.eps_Q.Q0", "L2.eps_Q.Q1"] +
               ["L3.eps_Q.Q0", "L3.eps_Q.Q1"] +
               ["L4.eps_Q.Q0", "L4.eps_Q.Q1"])
    assert_equal(mod.state_names, desired)
    desired = [
        "loading.0->M0", "loading.1->M0", "loading.0->M1",
        "loading.1->M1",
        "loading.0->Q0", "loading.1->Q0", "loading.0->Q1",
        "loading.1->Q1",
        "L1.0->0", "L1.1->0", "L2.0->0", "L2.1->0",
        "L3.0->0", "L3.1->0", "L4.0->0", "L4.1->0",
        "L5.0->0", "L5.1->0", "L6.0->0", "L6.1->0",
        "L1.0->1", "L1.1->1", "L2.0->1", "L2.1->1",
        "L3.0->1", "L3.1->1", "L4.0->1", "L4.1->1",
        "L5.0->1", "L5.1->1", "L6.0->1", "L6.1->1",
        "fb(0).cov.chol[1,1]", "fb(0).cov.chol[2,1]", "fb(0).cov.chol[2,2]",
        "sigma2.M0", "sigma2.M1", "sigma2.Q0", "sigma2.Q1"]
    assert_equal(mod.param_names, desired)

    # Test fixed elements of state space representation
    assert_allclose(mod["obs_intercept"], 0)

    assert_allclose(mod["design", 2:, 12:14], np.eye(2))
    assert_allclose(mod["design", 2:, 14:16], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 16:18], 3 * np.eye(2))
    assert_allclose(mod["design", 2:, 18:20], 2 * np.eye(2))
    assert_allclose(mod["design", 2:, 20:22], np.eye(2))
    assert_allclose(np.sum(mod["design"]), 18)

    assert_allclose(mod["obs_cov"], 0)
    assert_allclose(mod["state_intercept"], 0)

    assert_allclose(mod["transition", 2:12, :10], np.eye(10))
    assert_allclose(mod["transition", 14:22, 12:20], np.eye(2 * 4))
    assert_allclose(np.sum(mod["transition"]), 18)

    assert_allclose(mod["selection", :2, :2], np.eye(2))
    assert_allclose(mod["selection", 12:14, 2:4], np.eye(2))
    assert_allclose(np.sum(mod["selection"]), 4)

    assert_allclose(mod["state_cov"], 0)

    # Test parameter entry
    mod.update(np.arange(mod.k_params) + 2)

    # -> obs_intercept
    assert_allclose(mod["obs_intercept"], 0)

    # -> design
    desired = np.array([
        [2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [4., 5., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [6., 7., 12, 14, 18, 21, 12, 14, 6., 7., 0., 0.],
        [8., 9., 16, 18, 24, 27, 16, 18, 8., 9., 0., 0.]])
    assert_allclose(mod["design", :, :12], desired)
    desired = np.array([
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 0., 2., 0., 3., 0., 2., 0., 1., 0.],
        [0., 1., 0., 2., 0., 3., 0., 2., 0., 1.]])
    assert_allclose(mod["design", :, 12:], desired)

    # -> obs_cov
    assert_allclose(mod["obs_cov"], np.diag([37, 38, 0, 0]))

    # -> state_intercept
    assert_allclose(mod["state_intercept"], 0)

    # -> transition
    desired = np.array([
        [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
        [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", :12, :12], desired)
    assert_allclose(mod["transition", :12, 12:], 0)
    assert_allclose(mod["transition", 12:, :12], 0)
    desired = np.array([
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.]])
    assert_allclose(mod["transition", 12:, 12:], desired)

    # -> selection
    assert_allclose(np.sum(mod["selection"]), 4)

    # -> state_cov
    L = np.array([[34., 0],
                  [35., 36.]])
    desired = np.array([[0., 0., 0., 0.],
                        [0., 0., 0., 0.],
                        [0., 0., 39, 0.],
                        [0., 0., 0., 40]])
    desired[:2, :2] = np.dot(L, L.T)
    assert_allclose(mod["state_cov"], desired)


def test_invalid_model_specification():
    # Test for errors that can be raised for invalid model specifications
    # during __init__
    dta = np.zeros((10, 2))
    dta[0] = 1.
    dta_pd = pd.DataFrame(dta)
    dta_period_W = pd.DataFrame(
        dta, index=pd.period_range(start="2000", periods=10, freq="W"))
    dta_date_W = pd.DataFrame(
        dta, index=pd.date_range(start="2000", periods=10, freq="W"))
    dta_period_M = pd.DataFrame(
        dta, index=pd.period_range(start="2000", periods=10, freq="M"))
    dta_date_M = pd.DataFrame(
