import functools as ft
import warnings

from collections import defaultdict
from typing import Optional

import numpy as np
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
import pytest
import statsmodels.api as sm

from numpy.testing import assert_allclose
from scipy import linalg

from pymc_extras.statespace import structural as st
from pymc_extras.statespace.utils.constants import (
    ALL_STATE_AUX_DIM,
    ALL_STATE_DIM,
    AR_PARAM_DIM,
    OBS_STATE_AUX_DIM,
    OBS_STATE_DIM,
    SHOCK_AUX_DIM,
    SHOCK_DIM,
    SHORT_NAME_TO_LONG,
)
from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
    rng,
)
from tests.statespace.utilities.test_helpers import (
    assert_pattern_repeats,
    simulate_from_numpy_model,
    unpack_symbolic_matrices_with_params,
)

floatX = pytensor.config.floatX
ATOL = 1e-8 if floatX.endswith("64") else 1e-4
RTOL = 0 if floatX.endswith("64") else 1e-6


def _assert_all_statespace_matrices_match(mod, params, sm_mod):
    x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params)

    sm_x0, sm_H0, sm_P0 = sm_mod.initialization()

    if len(x0) > 0:
        assert_allclose(x0, sm_x0)

    for name, matrix in zip(["T", "R", "Z", "Q"], [T, R, Z, Q]):
        long_name = SHORT_NAME_TO_LONG[name]
        if np.any([x == 0 for x in matrix.shape]):
            continue
        assert_allclose(
            sm_mod.ssm[long_name],
            matrix,
            err_msg=f"matrix {name} does not match statsmodels",
            atol=ATOL,
            rtol=RTOL,
        )


def _assert_coord_shapes_match_matrices(mod, params):
    if "initial_state_cov" not in params:
        params["initial_state_cov"] = np.eye(mod.k_states)

    x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params)

    n_states = len(mod.coords[ALL_STATE_DIM])

    # There will always be one shock dimension -- dummies are inserted into fully deterministic models to avoid errors
    # in the state space representation.
    n_shocks = max(1, len(mod.coords[SHOCK_DIM]))
    n_obs = len(mod.coords[OBS_STATE_DIM])

    assert x0.shape[-1:] == (
        n_states,
    ), f"x0 expected to have shape (n_states, ), found {x0.shape[-1:]}"
    assert P0.shape[-2:] == (
        n_states,
        n_states,
    ), f"P0 expected to have shape (n_states, n_states), found {P0.shape[-2:]}"
    assert c.shape[-1:] == (
        n_states,
    ), f"c expected to have shape (n_states, ), found {c.shape[-1:]}"
    assert d.shape[-1:] == (n_obs,), f"d expected to have shape (n_obs, ), found {d.shape[-1:]}"
    assert T.shape[-2:] == (
        n_states,
        n_states,
    ), f"T expected to have shape (n_states, n_states), found {T.shape[-2:]}"
    assert Z.shape[-2:] == (
        n_obs,
        n_states,
    ), f"Z expected to have shape (n_obs, n_states), found {Z.shape[-2:]}"
    assert R.shape[-2:] == (
        n_states,
        n_shocks,
    ), f"R expected to have shape (n_states, n_shocks), found {R.shape[-2:]}"
    assert H.shape[-2:] == (
        n_obs,
        n_obs,
    ), f"H expected to have shape (n_obs, n_obs), found {H.shape[-2:]}"
    assert Q.shape[-2:] == (
        n_shocks,
        n_shocks,
    ), f"Q expected to have shape (n_shocks, n_shocks), found {Q.shape[-2:]}"


def _assert_basic_coords_correct(mod):
    assert mod.coords[ALL_STATE_DIM] == mod.state_names
    assert mod.coords[ALL_STATE_AUX_DIM] == mod.state_names
    assert mod.coords[SHOCK_DIM] == mod.shock_names
    assert mod.coords[SHOCK_AUX_DIM] == mod.shock_names
    assert mod.coords[OBS_STATE_DIM] == ["data"]
    assert mod.coords[OBS_STATE_AUX_DIM] == ["data"]


def _assert_keys_match(test_dict, expected_dict):
    expected_keys = list(expected_dict.keys())
    param_keys = list(test_dict.keys())
    key_diff = set(expected_keys) - set(param_keys)
    assert len(key_diff) == 0, f'{", ".join(key_diff)} were not found in the test_dict keys.'

    key_diff = set(param_keys) - set(expected_keys)
    assert (
        len(key_diff) == 0
    ), f'{", ".join(key_diff)} were keys of the tests_dict not in expected_dict.'


def _assert_param_dims_correct(param_dims, expected_dims):
    if len(expected_dims) == 0 and len(param_dims) == 0:
        return

    _assert_keys_match(param_dims, expected_dims)
    for param, dims in expected_dims.items():
        assert dims == param_dims[param], f"dims for parameter {param} do not match"


def _assert_coords_correct(coords, expected_coords):
    if len(coords) == 0 and len(expected_coords) == 0:
        return

    _assert_keys_match(coords, expected_coords)
    for dim, labels in expected_coords.items():
        assert labels == coords[dim], f"labels on dimension {dim} do not match"


def _assert_params_info_correct(param_info, coords, param_dims):
    for param in param_info.keys():
        info = param_info[param]

        dims = info["dims"]
        labels = [coords[dim] for dim in dims] if dims is not None else None
        if labels is not None:
            assert param in param_dims.keys()
            inferred_dims = param_dims[param]
        else:
            inferred_dims = None

        shape = tuple(len(label) for label in labels) if labels is not None else ()

        assert info["shape"] == shape
        assert dims == inferred_dims


def create_structural_model_and_equivalent_statsmodel(
    rng,
    level: bool | None = False,
    trend: bool | None = False,
    seasonal: int | None = None,
    freq_seasonal: list[dict] | None = None,
    cycle: bool = False,
    autoregressive: int | None = None,
    exog: np.ndarray | None = None,
    irregular: bool | None = False,
    stochastic_level: bool | None = True,
    stochastic_trend: bool | None = False,
    stochastic_seasonal: bool | None = True,
    stochastic_freq_seasonal: list[bool] | None = None,
    stochastic_cycle: bool | None = False,
    damped_cycle: bool | None = False,
):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mod = ft.partial(
            sm.tsa.UnobservedComponents,
            level=level,
            trend=trend,
            seasonal=seasonal,
            freq_seasonal=freq_seasonal,
            cycle=cycle,
            autoregressive=autoregressive,
            exog=exog,
            irregular=irregular,
            stochastic_level=stochastic_level,
            stochastic_trend=stochastic_trend,
            stochastic_seasonal=stochastic_seasonal,
            stochastic_freq_seasonal=stochastic_freq_seasonal,
            stochastic_cycle=stochastic_cycle,
            damped_cycle=damped_cycle,
            mle_regression=False,
        )

    params = {}
    sm_params = {}
    sm_init = {}
    expected_param_dims = defaultdict(tuple)
    expected_coords = defaultdict(list)
    expected_param_dims["P0"] += ("state", "state_aux")

    default_states = [
        ALL_STATE_DIM,
        ALL_STATE_AUX_DIM,
        OBS_STATE_DIM,
        OBS_STATE_AUX_DIM,
        SHOCK_DIM,
        SHOCK_AUX_DIM,
    ]
    default_values = [[], [], ["data"], ["data"], [], []]
    for dim, value in zip(default_states, default_values):
        expected_coords[dim] += value

    components = []

    if irregular:
        sigma2 = np.abs(rng.normal()).astype(floatX).item()
        params["sigma_irregular"] = np.sqrt(sigma2)
        sm_params["sigma2.irregular"] = sigma2

        comp = st.MeasurementError("irregular")
        components.append(comp)

    level_trend_order = [0, 0]
    level_trend_innov_order = [0, 0]

    if level:
        level_trend_order[0] = 1
        expected_coords["trend_state"] += [
            "level",
        ]
        expected_coords[ALL_STATE_DIM] += [
            "level",
        ]
        expected_coords[ALL_STATE_AUX_DIM] += [
            "level",
        ]
        if stochastic_level:
            level_trend_innov_order[0] = 1
            expected_coords["trend_shock"] += ["level"]
            expected_coords[SHOCK_DIM] += [
                "level",
            ]
            expected_coords[SHOCK_AUX_DIM] += [
                "level",
            ]

    if trend:
        level_trend_order[1] = 1
        expected_coords["trend_state"] += [
            "trend",
        ]
        expected_coords[ALL_STATE_DIM] += [
            "trend",
        ]
        expected_coords[ALL_STATE_AUX_DIM] += [
            "trend",
        ]

        if stochastic_trend:
            level_trend_innov_order[1] = 1
            expected_coords["trend_shock"] += ["trend"]
            expected_coords[SHOCK_DIM] += ["trend"]
            expected_coords[SHOCK_AUX_DIM] += ["trend"]

    if level or trend:
        expected_param_dims["initial_trend"] += ("trend_state",)
        level_value = np.where(
            level_trend_order,
            rng.normal(
                size=2,
            ).astype(floatX),
            np.zeros(2, dtype=floatX),
        )
        sigma_level_value2 = np.abs(rng.normal(size=(2,)))[
            np.array(level_trend_innov_order, dtype="bool")
        ]
        max_order = np.flatnonzero(level_value)[-1].item() + 1
        level_trend_order = level_trend_order[:max_order]

        params["initial_trend"] = level_value[:max_order]
        sm_init["level"] = level_value[0]
        sm_init["trend"] = level_value[1]

        if sum(level_trend_innov_order) > 0:
            expected_param_dims["sigma_trend"] += ("trend_shock",)
            params["sigma_trend"] = np.sqrt(sigma_level_value2)

        sigma_level_value = sigma_level_value2.tolist()
        if stochastic_level:
            sigma = sigma_level_value.pop(0)
            sm_params["sigma2.level"] = sigma
        if stochastic_trend:
            sigma = sigma_level_value.pop(0)
            sm_params["sigma2.trend"] = sigma

        comp = st.LevelTrendComponent(
            name="level", order=level_trend_order, innovations_order=level_trend_innov_order
        )
        components.append(comp)

    if seasonal is not None:
        state_names = [f"seasonal_{i}" for i in range(seasonal)][1:]
        seasonal_coefs = rng.normal(size=(seasonal - 1,)).astype(floatX)
        params["seasonal_coefs"] = seasonal_coefs
        expected_param_dims["seasonal_coefs"] += ("seasonal_state",)

        expected_coords["seasonal_state"] += tuple(state_names)
        expected_coords[ALL_STATE_DIM] += state_names
        expected_coords[ALL_STATE_AUX_DIM] += state_names

        seasonal_dict = {
            "seasonal" if i == 0 else f"seasonal.L{i}": c for i, c in enumerate(seasonal_coefs)
        }
        sm_init.update(seasonal_dict)

        if stochastic_seasonal:
            sigma2 = np.abs(rng.normal()).astype(floatX)
            params["sigma_seasonal"] = np.sqrt(sigma2)
            sm_params["sigma2.seasonal"] = sigma2
            expected_coords[SHOCK_DIM] += [
                "seasonal",
            ]
            expected_coords[SHOCK_AUX_DIM] += [
                "seasonal",
            ]

        comp = st.TimeSeasonality(
            name="seasonal", season_length=seasonal, innovations=stochastic_seasonal
        )
        components.append(comp)

    if freq_seasonal is not None:
        state_count = 0
        for d, has_innov in zip(freq_seasonal, stochastic_freq_seasonal):
            n = d["harmonics"]
            s = d["period"]
            last_state_not_identified = (s / n) == 2.0
            n_states = 2 * n - int(last_state_not_identified)
            state_names = [f"seasonal_{s}_{f}_{i}" for i in range(n) for f in ["Cos", "Sin"]]

            seasonal_params = rng.normal(size=n_states).astype(floatX)

            params[f"seasonal_{s}"] = seasonal_params
            expected_param_dims[f"seasonal_{s}"] += (f"seasonal_{s}_state",)
            expected_coords[ALL_STATE_DIM] += state_names
            expected_coords[ALL_STATE_AUX_DIM] += state_names
            expected_coords[f"seasonal_{s}_state"] += (
                tuple(state_names[:-1]) if last_state_not_identified else tuple(state_names)
            )

            for param in seasonal_params:
                sm_init[f"freq_seasonal.{state_count}"] = param
                state_count += 1
            if last_state_not_identified:
                sm_init[f"freq_seasonal.{state_count}"] = 0.0
                state_count += 1

            if has_innov:
                sigma2 = np.abs(rng.normal()).astype(floatX)
                params[f"sigma_seasonal_{s}"] = np.sqrt(sigma2)
                sm_params[f"sigma2.freq_seasonal_{s}({n})"] = sigma2
                expected_coords[SHOCK_DIM] += state_names
                expected_coords[SHOCK_AUX_DIM] += state_names

            comp = st.FrequencySeasonality(
                name=f"seasonal_{s}", season_length=s, n=n, innovations=has_innov
            )
            components.append(comp)

    if cycle:
        cycle_length = np.random.choice(np.arange(2, 12)).astype(floatX)

        # Statsmodels takes the frequency not the cycle length, so convert it.
        sm_params["frequency.cycle"] = 2.0 * np.pi / cycle_length
        params["cycle_length"] = cycle_length

        init_cycle = rng.normal(size=(2,)).astype(floatX)
        params["cycle"] = init_cycle
        expected_param_dims["cycle"] += ("cycle_state",)

        state_names = ["cycle_Cos", "cycle_Sin"]
        expected_coords["cycle_state"] += state_names
        expected_coords[ALL_STATE_DIM] += state_names
        expected_coords[ALL_STATE_AUX_DIM] += state_names

        sm_init["cycle"] = init_cycle[0]
        sm_init["cycle.auxilliary"] = init_cycle[1]

        if stochastic_cycle:
            sigma2 = np.abs(rng.normal()).astype(floatX)
            params["sigma_cycle"] = np.sqrt(sigma2)
            expected_coords[SHOCK_DIM] += state_names
            expected_coords[SHOCK_AUX_DIM] += state_names

            sm_params["sigma2.cycle"] = sigma2

        if damped_cycle:
            rho = rng.beta(1, 1)
            params["cycle_dampening_factor"] = rho
            sm_params["damping.cycle"] = rho

        comp = st.CycleComponent(
            name="cycle",
            dampen=damped_cycle,
            innovations=stochastic_cycle,
            estimate_cycle_length=True,
        )

        components.append(comp)

    if autoregressive is not None:
        ar_names = [f"L{i+1}.data" for i in range(autoregressive)]
        ar_params = rng.normal(size=(autoregressive,)).astype(floatX)
        if autoregressive == 1:
            ar_params = ar_params.item()
        sigma2 = np.abs(rng.normal()).astype(floatX)

        params["ar_params"] = ar_params
        params["sigma_ar"] = np.sqrt(sigma2)
        expected_param_dims["ar_params"] += (AR_PARAM_DIM,)
        expected_coords[AR_PARAM_DIM] += tuple(list(range(1, autoregressive + 1)))
        expected_coords[ALL_STATE_DIM] += ar_names
        expected_coords[ALL_STATE_AUX_DIM] += ar_names
        expected_coords[SHOCK_DIM] += ["ar_innovation"]
        expected_coords[SHOCK_AUX_DIM] += ["ar_innovation"]

        sm_params["sigma2.ar"] = sigma2
        for i, rho in enumerate(ar_params):
            sm_init[f"ar.L{i+1}"] = 0
            sm_params[f"ar.L{i+1}"] = rho

        comp = st.AutoregressiveComponent(name="ar", order=autoregressive)
        components.append(comp)

    if exog is not None:
        names = [f"x{i + 1}" for i in range(exog.shape[1])]
        betas = rng.normal(size=(exog.shape[1],)).astype(floatX)
        params["beta_exog"] = betas
        params["data_exog"] = exog
        expected_param_dims["beta_exog"] += ("exog_state",)
        expected_param_dims["data_exog"] += ("time", "exog_data")

        expected_coords["exog_state"] += tuple(names)

        for i, beta in enumerate(betas):
            sm_params[f"beta.x{i + 1}"] = beta
            sm_init[f"beta.x{i+1}"] = beta
        comp = st.RegressionComponent(name="exog", state_names=names)
        components.append(comp)

    st_mod = components.pop(0)
    for comp in components:
        st_mod += comp
    return mod, st_mod, params, sm_params, sm_init, expected_param_dims, expected_coords


@pytest.mark.parametrize(
    "level, trend, stochastic_level, stochastic_trend, irregular",
    [
        (False, False, False, False, True),
        (True, True, True, True, True),
        (True, True, False, True, False),
    ],
)
@pytest.mark.parametrize("autoregressive", [None, 3])
@pytest.mark.parametrize("seasonal, stochastic_seasonal", [(None, False), (12, False), (12, True)])
@pytest.mark.parametrize(
    "freq_seasonal, stochastic_freq_seasonal",
    [
        (None, None),
        ([{"period": 12, "harmonics": 2}], [False]),
        ([{"period": 12, "harmonics": 6}], [True]),
    ],
)
@pytest.mark.parametrize(
    "cycle, damped_cycle, stochastic_cycle",
    [(False, False, False), (True, False, True), (True, True, True)],
)
@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.ConvergenceWarning")
@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.SpecificationWarning")
def test_structural_model_against_statsmodels(
    level,
    trend,
    stochastic_level,
    stochastic_trend,
    irregular,
    autoregressive,
    seasonal,
    stochastic_seasonal,
    freq_seasonal,
    stochastic_freq_seasonal,
    cycle,
    damped_cycle,
    stochastic_cycle,
    rng,
):
    retvals = create_structural_model_and_equivalent_statsmodel(
        rng,
        level=level,
        trend=trend,
        seasonal=seasonal,
        freq_seasonal=freq_seasonal,
        cycle=cycle,
        damped_cycle=damped_cycle,
        autoregressive=autoregressive,
        irregular=irregular,
        stochastic_level=stochastic_level,
        stochastic_trend=stochastic_trend,
        stochastic_seasonal=stochastic_seasonal,
        stochastic_freq_seasonal=stochastic_freq_seasonal,
        stochastic_cycle=stochastic_cycle,
    )
    f_sm_mod, mod, params, sm_params, sm_init, expected_dims, expected_coords = retvals

    data = rng.normal(size=(100,)).astype(floatX)
    sm_mod = f_sm_mod(data)

    if len(sm_init) > 0:
        init_array = np.concatenate(
            [np.atleast_1d(sm_init[k]).ravel() for k in sm_mod.state_names if k != "dummy"]
        )
        sm_mod.initialize_known(init_array, np.eye(sm_mod.k_states))
    else:
        sm_mod.initialize_default()

    if len(sm_params) > 0:
        param_array = np.concatenate(
            [np.atleast_1d(sm_params[k]).ravel() for k in sm_mod.param_names]
        )
        sm_mod.update(param_array, transformed=True)

    _assert_all_statespace_matrices_match(mod, params, sm_mod)

    built_model = mod.build(verbose=False)

    _assert_coord_shapes_match_matrices(built_model, params)
    _assert_param_dims_correct(built_model.param_dims, expected_dims)
    _assert_coords_correct(built_model.coords, expected_coords)
    _assert_params_info_correct(built_model.param_info, built_model.coords, built_model.param_dims)


def test_level_trend_model(rng):
    mod = st.LevelTrendComponent(order=2, innovations_order=0)
    params = {"initial_trend": [0.0, 1.0]}
    x, y = simulate_from_numpy_model(mod, rng, params)

    assert_allclose(np.diff(y), 1, atol=ATOL, rtol=RTOL)

    # Check coords
    mod = mod.build(verbose=False)
    _assert_basic_coords_correct(mod)
    assert mod.coords["trend_state"] == ["level", "trend"]


def test_measurement_error(rng):
    mod = st.MeasurementError("obs") + st.LevelTrendComponent(order=2)
    mod = mod.build(verbose=False)

    _assert_basic_coords_correct(mod)
    assert "sigma_obs" in mod.param_names


@pytest.mark.parametrize("order", [1, 2, [1, 0, 1]], ids=["AR1", "AR2", "AR(1,0,1)"])
def test_autoregressive_model(order, rng):
    ar = st.AutoregressiveComponent(order=order)
    params = {
        "ar_params": np.full((sum(ar.order),), 0.5, dtype=floatX),
        "sigma_ar": 0.0,
    }

    x, y = simulate_from_numpy_model(ar, rng, params, steps=100)

    # Check coords
    ar.build(verbose=False)
    _assert_basic_coords_correct(ar)
    lags = np.arange(len(order) if isinstance(order, list) else order, dtype="int") + 1
    if isinstance(order, list):
        lags = lags[np.flatnonzero(order)]
    assert_allclose(ar.coords["ar_lag"], lags)


@pytest.mark.parametrize("s", [10, 25, 50])
@pytest.mark.parametrize("innovations", [True, False])
def test_time_seasonality(s, innovations, rng):
    def random_word(rng):
        return "".join(rng.choice(list("abcdefghijklmnopqrstuvwxyz")) for _ in range(5))

    state_names = [random_word(rng) for _ in range(s)]
    mod = st.TimeSeasonality(
        season_length=s, innovations=innovations, name="season", state_names=state_names
    )
    x0 = np.zeros(mod.k_states, dtype=floatX)
    x0[0] = 1

    params = {"season_coefs": x0}
    if mod.innovations:
        params["sigma_season"] = 0.0

    x, y = simulate_from_numpy_model(mod, rng, params)
    y = y.ravel()
    if not innovations:
        assert_pattern_repeats(y, s, atol=ATOL, rtol=RTOL)

    # Check coords
    mod.build(verbose=False)
    _assert_basic_coords_correct(mod)
    assert mod.coords["season_state"] == state_names[1:]


def get_shift_factor(s):
    s_str = str(s)
    if "." not in s_str:
        return 1
    _, decimal = s_str.split(".")
    return 10 ** len(decimal)


@pytest.mark.parametrize("n", [*np.arange(1, 6, dtype="int").tolist(), None])
@pytest.mark.parametrize("s", [5, 10, 25, 25.2])
def test_frequency_seasonality(n, s, rng):
    mod = st.FrequencySeasonality(season_length=s, n=n, name="season")
    x0 = rng.normal(size=mod.n_coefs).astype(floatX)
    params = {"season": x0, "sigma_season": 0.0}
    k = get_shift_factor(s)
    T = int(s * k)

    x, y = simulate_from_numpy_model(mod, rng, params, steps=2 * T)
    assert_pattern_repeats(y, T, atol=ATOL, rtol=RTOL)

    # Check coords
    mod.build(verbose=False)
    _assert_basic_coords_correct(mod)
    if n is None:
        n = int(s // 2)
    states = [f"season_{f}_{i}" for i in range(n) for f in ["Cos", "Sin"]]

    # Remove the last state when the model is completely saturated
    if s / n == 2.0:
        states.pop()
    assert mod.coords["season_state"] == states


cycle_test_vals = zip([None, None, 3, 5, 10], [False, True, True, False, False])


def test_cycle_component_deterministic(rng):
    cycle = st.CycleComponent(
        name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False
    )
    params = {"cycle": np.array([1.0, 1.0], dtype=floatX)}
    x, y = simulate_from_numpy_model(cycle, rng, params, steps=12 * 12)

    assert_pattern_repeats(y, 12, atol=ATOL, rtol=RTOL)


def test_cycle_component_with_dampening(rng):
    cycle = st.CycleComponent(
        name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False, dampen=True
    )
    params = {"cycle": np.array([10.0, 10.0], dtype=floatX), "cycle_dampening_factor": 0.75}
    x, y = simulate_from_numpy_model(cycle, rng, params, steps=100)

    # Check that the cycle dampens to zero over time
    assert_allclose(y[-1], 0.0, atol=ATOL, rtol=RTOL)


def test_cycle_component_with_innovations_and_cycle_length(rng):
    cycle = st.CycleComponent(
        name="cycle", estimate_cycle_length=True, innovations=True, dampen=True
    )
    params = {
        "cycle": np.array([1.0, 1.0], dtype=floatX),
        "cycle_length": 12.0,
        "cycle_dampening_factor": 0.95,
        "sigma_cycle": 1.0,
    }

    x, y = simulate_from_numpy_model(cycle, rng, params)

    cycle.build(verbose=False)
    _assert_basic_coords_correct(cycle)


def test_exogenous_component(rng):
    data = rng.normal(size=(100, 2)).astype(floatX)
    mod = st.RegressionComponent(state_names=["feature_1", "feature_2"], name="exog")

    params = {"beta_exog": np.array([1.0, 2.0], dtype=floatX)}
    exog_data = {"data_exog": data}
    x, y = simulate_from_numpy_model(mod, rng, params, exog_data)

    # Check that the generated data is just a linear regression
    assert_allclose(y, data @ params["beta_exog"], atol=ATOL, rtol=RTOL)

    mod.build(verbose=False)
    _assert_basic_coords_correct(mod)
    assert mod.coords["exog_state"] == ["feature_1", "feature_2"]


def test_adding_exogenous_component(rng):
    data = rng.normal(size=(100, 2)).astype(floatX)
    reg = st.RegressionComponent(state_names=["a", "b"], name="exog")
    ll = st.LevelTrendComponent(name="level")

    seasonal = st.FrequencySeasonality(name="annual", season_length=12, n=4)
    mod = reg + ll + seasonal

    assert mod.ssm["design"].eval({"data_exog": data}).shape == (100, 1, 2 + 2 + 8)
    assert_allclose(mod.ssm["design", 5, 0, :2].eval({"data_exog": data}), data[5])


def test_add_components():
    ll = st.LevelTrendComponent(order=2)
    se = st.TimeSeasonality(name="seasonal", season_length=12)
    mod = ll + se

    ll_params = {
        "initial_trend": np.zeros(2, dtype=floatX),
        "sigma_trend": np.ones(2, dtype=floatX),
    }
    se_params = {
        "seasonal_coefs": np.ones(11, dtype=floatX),
        "sigma_seasonal": 1.0,
    }
    all_params = ll_params.copy()
    all_params.update(se_params)

    (ll_x0, ll_P0, ll_c, ll_d, ll_T, ll_Z, ll_R, ll_H, ll_Q) = unpack_symbolic_matrices_with_params(
        ll, ll_params
    )
    (se_x0, se_P0, se_c, se_d, se_T, se_Z, se_R, se_H, se_Q) = unpack_symbolic_matrices_with_params(
        se, se_params
    )
    x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, all_params)

    for property in ["param_names", "shock_names", "param_info", "coords", "param_dims"]:
        assert [x in getattr(mod, property) for x in getattr(ll, property)]
        assert [x in getattr(mod, property) for x in getattr(se, property)]

    ll_mats = [ll_T, ll_R, ll_Q]
    se_mats = [se_T, se_R, se_Q]
    all_mats = [T, R, Q]

    for ll_mat, se_mat, all_mat in zip(ll_mats, se_mats, all_mats):
        assert_allclose(all_mat, linalg.block_diag(ll_mat, se_mat), atol=ATOL, rtol=RTOL)

    ll_mats = [ll_x0, ll_c, ll_Z]
    se_mats = [se_x0, se_c, se_Z]
    all_mats = [x0, c, Z]
    axes = [0, 0, 1]

    for ll_mat, se_mat, all_mat, axis in zip(ll_mats, se_mats, all_mats, axes):
        assert_allclose(all_mat, np.concatenate([ll_mat, se_mat], axis=axis), atol=ATOL, rtol=RTOL)


def test_filter_scans_time_varying_design_matrix(rng):
    time_idx = pd.date_range(start="2000-01-01", freq="D", periods=100)
    data = pd.DataFrame(rng.normal(size=(100, 2)), columns=["a", "b"], index=time_idx)

    y = pd.DataFrame(rng.normal(size=(100, 1)), columns=["data"], index=time_idx)

    reg = st.RegressionComponent(state_names=["a", "b"], name="exog")
    mod = reg.build(verbose=False)

    with pm.Model(coords=mod.coords) as m:
        data_exog = pm.Data("data_exog", data.values)

        x0 = pm.Normal("x0", dims=["state"])
        P0 = pm.Deterministic("P0", pt.eye(mod.k_states), dims=["state", "state_aux"])
        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])

        mod.build_statespace_graph(y)
        x0, P0, c, d, T, Z, R, H, Q = mod.unpack_statespace()
        pm.Deterministic("Z", Z)

        prior = pm.sample_prior_predictive(draws=10)

    prior_Z = prior.prior.Z.values
    assert prior_Z.shape == (1, 10, 100, 1, 2)
    assert_allclose(prior_Z[0, :, :, 0, :], data.values[None].repeat(10, axis=0))


@pytest.mark.skipif(floatX.endswith("32"), reason="Prior covariance not PSD at half-precision")
def test_extract_components_from_idata(rng):
    time_idx = pd.date_range(start="2000-01-01", freq="D", periods=100)
    data = pd.DataFrame(rng.normal(size=(100, 2)), columns=["a", "b"], index=time_idx)

    y = pd.DataFrame(rng.normal(size=(100, 1)), columns=["data"], index=time_idx)

    ll = st.LevelTrendComponent()
    season = st.FrequencySeasonality(name="seasonal", season_length=12, n=2, innovations=False)
    reg = st.RegressionComponent(state_names=["a", "b"], name="exog")
    me = st.MeasurementError("obs")
    mod = (ll + season + reg + me).build(verbose=False)

    with pm.Model(coords=mod.coords) as m:
        data_exog = pm.Data("data_exog", data.values)

        x0 = pm.Normal("x0", dims=["state"])
        P0 = pm.Deterministic("P0", pt.eye(mod.k_states), dims=["state", "state_aux"])
        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
        seasonal_coefs = pm.Normal("seasonal", dims=["seasonal_state"])
        sigma_obs = pm.Exponential("sigma_obs", 1)

        mod.build_statespace_graph(y)

        x0, P0, c, d, T, Z, R, H, Q = mod.unpack_statespace()
        prior = pm.sample_prior_predictive(draws=10)

    filter_prior = mod.sample_conditional_prior(prior)
    comp_prior = mod.extract_components_from_idata(filter_prior)
    comp_states = comp_prior.filtered_prior.coords["state"].values
    expected_states = ["LevelTrend[level]", "LevelTrend[trend]", "seasonal", "exog[a]", "exog[b]"]
    missing = set(comp_states) - set(expected_states)

    assert len(missing) == 0, missing
