from functools import partial

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

from numpy.testing import assert_allclose

from pymc_extras.statespace.core.statespace import FILTER_FACTORY, PyMCStateSpace
from pymc_extras.statespace.models import structural as st
from pymc_extras.statespace.models.utilities import make_default_coords
from pymc_extras.statespace.utils.constants import (
    FILTER_OUTPUT_NAMES,
    MATRIX_NAMES,
    SMOOTHER_OUTPUT_NAMES,
)
from tests.statespace.utilities.shared_fixtures import (
    rng,
)
from tests.statespace.utilities.test_helpers import (
    fast_eval,
    load_nile_test_data,
    make_test_inputs,
)

floatX = pytensor.config.floatX
nile = load_nile_test_data()
ALL_SAMPLE_OUTPUTS = MATRIX_NAMES + FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES


def make_statespace_mod(k_endog, k_states, k_posdef, filter_type, verbose=False, data_info=None):
    class StateSpace(PyMCStateSpace):
        def make_symbolic_graph(self):
            pass

        @property
        def data_info(self):
            return data_info

    ss = StateSpace(
        k_states=k_states,
        k_endog=k_endog,
        k_posdef=k_posdef,
        filter_type=filter_type,
        verbose=verbose,
    )
    ss._needs_exog_data = data_info is not None
    ss._exog_names = list(data_info.keys()) if data_info is not None else []

    return ss


@pytest.fixture(scope="session")
def ss_mod():
    class StateSpace(PyMCStateSpace):
        @property
        def param_names(self):
            return ["rho", "zeta"]

        @property
        def state_names(self):
            return ["a", "b"]

        @property
        def observed_states(self):
            return ["a"]

        @property
        def shock_names(self):
            return ["a"]

        @property
        def coords(self):
            return make_default_coords(self)

        def make_symbolic_graph(self):
            rho = self.make_and_register_variable("rho", ())
            zeta = self.make_and_register_variable("zeta", ())
            self.ssm["transition", 0, 0] = rho
            self.ssm["transition", 1, 0] = zeta

    Z = np.array([[1.0, 0.0]], dtype=floatX)
    R = np.array([[1.0], [0.0]], dtype=floatX)
    H = np.array([[0.1]], dtype=floatX)
    Q = np.array([[0.8]], dtype=floatX)
    P0 = np.eye(2, dtype=floatX) * 1e6

    ss_mod = StateSpace(
        k_endog=nile.shape[1], k_states=2, k_posdef=1, filter_type="standard", verbose=False
    )
    for X, name in zip(
        [Z, R, H, Q, P0],
        ["design", "selection", "obs_cov", "state_cov", "initial_state_cov"],
    ):
        ss_mod.ssm[name] = X

    return ss_mod


@pytest.fixture(scope="session")
def pymc_mod(ss_mod):
    with pm.Model(coords=ss_mod.coords) as pymc_mod:
        rho = pm.Beta("rho", 1, 1)
        zeta = pm.Deterministic("zeta", 1 - rho)

        ss_mod.build_statespace_graph(data=nile, save_kalman_filter_outputs_in_idata=True)
        names = ["x0", "P0", "c", "d", "T", "Z", "R", "H", "Q"]
        for name, matrix in zip(names, ss_mod.unpack_statespace()):
            pm.Deterministic(name, matrix)

    return pymc_mod


@pytest.fixture(scope="session")
def ss_mod_no_exog(rng):
    ll = st.LevelTrendComponent(order=2, innovations_order=1)
    return ll.build()


@pytest.fixture(scope="session")
def ss_mod_no_exog_dt(rng):
    ll = st.LevelTrendComponent(order=2, innovations_order=1)
    return ll.build()


@pytest.fixture(scope="session")
def exog_ss_mod(rng):
    ll = st.LevelTrendComponent()
    reg = st.RegressionComponent(name="exog", state_names=["a", "b", "c"])
    mod = (ll + reg).build(verbose=False)

    return mod


@pytest.fixture(scope="session")
def exog_pymc_mod(exog_ss_mod, rng):
    y = rng.normal(size=(100, 1)).astype(floatX)
    X = rng.normal(size=(100, 3)).astype(floatX)

    with pm.Model(coords=exog_ss_mod.coords) as m:
        exog_data = pm.Data("data_exog", X)
        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
        P0_sigma = pm.Exponential("P0_sigma", 1)
        P0 = pm.Deterministic(
            "P0", pt.eye(exog_ss_mod.k_states) * P0_sigma, dims=["state", "state_aux"]
        )
        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])

        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
        exog_ss_mod.build_statespace_graph(y, save_kalman_filter_outputs_in_idata=True)

    return m


@pytest.fixture(scope="session")
def pymc_mod_no_exog(ss_mod_no_exog, rng):
    y = pd.DataFrame(rng.normal(size=(100, 1)).astype(floatX), columns=["y"])

    with pm.Model(coords=ss_mod_no_exog.coords) as m:
        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
        P0_sigma = pm.Exponential("P0_sigma", 1)
        P0 = pm.Deterministic(
            "P0", pt.eye(ss_mod_no_exog.k_states) * P0_sigma, dims=["state", "state_aux"]
        )
        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
        ss_mod_no_exog.build_statespace_graph(y)

    return m


@pytest.fixture(scope="session")
def pymc_mod_no_exog_dt(ss_mod_no_exog_dt, rng):
    y = pd.DataFrame(
        rng.normal(size=(100, 1)).astype(floatX),
        columns=["y"],
        index=pd.date_range("2020-01-01", periods=100, freq="D"),
    )

    with pm.Model(coords=ss_mod_no_exog_dt.coords) as m:
        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
        P0_sigma = pm.Exponential("P0_sigma", 1)
        P0 = pm.Deterministic(
            "P0", pt.eye(ss_mod_no_exog_dt.k_states) * P0_sigma, dims=["state", "state_aux"]
        )
        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
        ss_mod_no_exog_dt.build_statespace_graph(y)

    return m


@pytest.fixture(scope="session")
def idata(pymc_mod, rng):
    with pymc_mod:
        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)

    idata.extend(idata_prior)
    return idata


@pytest.fixture(scope="session")
def idata_exog(exog_pymc_mod, rng):
    with exog_pymc_mod:
        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)
    idata.extend(idata_prior)
    return idata


@pytest.fixture(scope="session")
def idata_no_exog(pymc_mod_no_exog, rng):
    with pymc_mod_no_exog:
        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)
    idata.extend(idata_prior)
    return idata


@pytest.fixture(scope="session")
def idata_no_exog_dt(pymc_mod_no_exog_dt, rng):
    with pymc_mod_no_exog_dt:
        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)
    idata.extend(idata_prior)
    return idata


def test_invalid_filter_name_raises():
    msg = "The following are valid filter types: " + ", ".join(list(FILTER_FACTORY.keys()))
    with pytest.raises(NotImplementedError, match=msg):
        mod = make_statespace_mod(k_endog=1, k_states=5, k_posdef=1, filter_type="invalid_filter")


def test_unpack_before_insert_raises(rng):
    p, m, r, n = 2, 5, 1, 10
    data, *inputs = make_test_inputs(p, m, r, n, rng, missing_data=0)
    mod = make_statespace_mod(
        k_endog=p, k_states=m, k_posdef=r, filter_type="standard", verbose=False
    )

    msg = "Cannot unpack the complete statespace system until PyMC model variables have been inserted."
    with pytest.raises(ValueError, match=msg):
        outputs = mod.unpack_statespace()


def test_unpack_matrices(rng):
    p, m, r, n = 2, 5, 1, 10
    data, *inputs = make_test_inputs(p, m, r, n, rng, missing_data=0)
    mod = make_statespace_mod(
        k_endog=p, k_states=m, k_posdef=r, filter_type="standard", verbose=False
    )

    # mod is a dummy statespace, so there are no placeholders to worry about. Monkey patch subbed_ssm with the defaults
    mod.subbed_ssm = mod._unpack_statespace_with_placeholders()

    outputs = mod.unpack_statespace()
    for x, y in zip(inputs, outputs):
        assert_allclose(np.zeros_like(x), fast_eval(y))


def test_param_names_raises_on_base_class():
    mod = make_statespace_mod(
        k_endog=1, k_states=5, k_posdef=1, filter_type="standard", verbose=False
    )
    with pytest.raises(NotImplementedError):
        x = mod.param_names


def test_base_class_raises():
    with pytest.raises(NotImplementedError):
        mod = PyMCStateSpace(
            k_endog=1, k_states=5, k_posdef=1, filter_type="standard", verbose=False
        )


def test_update_raises_if_missing_variables(ss_mod):
    with pm.Model() as mod:
        rho = pm.Normal("rho")
        msg = "The following required model parameters were not found in the PyMC model: zeta"
        with pytest.raises(ValueError, match=msg):
            ss_mod._insert_random_variables()


def test_build_statespace_graph_warns_if_data_has_nans():
    # Breaks tests if it uses the session fixtures because we can't call build_statespace_graph over and over
    ss_mod = st.LevelTrendComponent(order=1, innovations_order=0).build(verbose=False)

    with pm.Model() as pymc_mod:
        initial_trend = pm.Normal("initial_trend", shape=(1,))
        P0 = pm.Deterministic("P0", pt.eye(1, dtype=floatX))
        with pytest.warns(pm.ImputationWarning):
            ss_mod.build_statespace_graph(
                data=np.full((10, 1), np.nan, dtype=floatX), register_data=False
            )


def test_build_statespace_graph_raises_if_data_has_missing_fill():
    # Breaks tests if it uses the session fixtures because we can't call build_statespace_graph over and over
    ss_mod = st.LevelTrendComponent(order=1, innovations_order=0).build(verbose=False)

    with pm.Model() as pymc_mod:
        initial_trend = pm.Normal("initial_trend", shape=(1,))
        P0 = pm.Deterministic("P0", pt.eye(1, dtype=floatX))
        with pytest.raises(ValueError, match="Provided data contains the value 1.0"):
            data = np.ones((10, 1), dtype=floatX)
            data[3] = np.nan
            ss_mod.build_statespace_graph(data=data, missing_fill_value=1.0, register_data=False)


def test_build_statespace_graph(pymc_mod):
    for name in [
        "filtered_state",
        "predicted_state",
        "predicted_covariance",
        "filtered_covariance",
    ]:
        assert name in [x.name for x in pymc_mod.deterministics]


def test_build_smoother_graph(ss_mod, pymc_mod):
    names = ["smoothed_state", "smoothed_covariance"]
    for name in names:
        assert name in [x.name for x in pymc_mod.deterministics]


@pytest.mark.parametrize("group", ["posterior", "prior"])
@pytest.mark.parametrize("matrix", ALL_SAMPLE_OUTPUTS)
def test_no_nans_in_sampling_output(group, matrix, idata):
    assert not np.any(np.isnan(idata[group][matrix].values))


@pytest.mark.parametrize("group", ["posterior", "prior"])
@pytest.mark.parametrize("kind", ["conditional", "unconditional"])
def test_sampling_methods(group, kind, ss_mod, idata, rng):
    f = getattr(ss_mod, f"sample_{kind}_{group}")
    test_idata = f(idata, random_seed=rng)

    if kind == "conditional":
        for output in ["filtered", "predicted", "smoothed"]:
            assert f"{output}_{group}" in test_idata
            assert not np.any(np.isnan(test_idata[f"{output}_{group}"].values))
            assert not np.any(np.isnan(test_idata[f"{output}_{group}_observed"].values))
    if kind == "unconditional":
        for output in ["latent", "observed"]:
            assert f"{group}_{output}" in test_idata
            assert not np.any(np.isnan(test_idata[f"{group}_{output}"].values))


def _make_time_idx(mod, use_datetime_index=True):
    if use_datetime_index:
        mod._fit_coords["time"] = nile.index
        time_idx = nile.index
    else:
        mod._fit_coords["time"] = nile.reset_index().index
        time_idx = pd.RangeIndex(start=0, stop=nile.shape[0], step=1)

    return time_idx


@pytest.mark.parametrize("use_datetime_index", [True, False])
def test_bad_forecast_arguments(use_datetime_index, caplog):
    ss_mod = make_statespace_mod(
        k_endog=1, k_posdef=1, k_states=2, filter_type="standard", verbose=False
    )

    # Not-fit model raises
    ss_mod._fit_coords = dict()
    with pytest.raises(ValueError, match="Has this model been fit?"):
        ss_mod._get_fit_time_index()

    time_idx = _make_time_idx(ss_mod, use_datetime_index)

    # Start value not in time index
    match = (
        "Datetime start must be in the data index used to fit the model"
        if use_datetime_index
        else "Integer start must be within the range of the data index used to fit the model."
    )
    with pytest.raises(ValueError, match=match):
        start = time_idx.shift(10)[-1] if use_datetime_index else time_idx[-1] + 11
        ss_mod._validate_forecast_args(time_index=time_idx, start=start, periods=10)

    # End value cannot be inferred
    with pytest.raises(ValueError, match="Must specify one of either periods or end"):
        start = time_idx[-1]
        ss_mod._validate_forecast_args(time_index=time_idx, start=start)

    # Unnecessary args warn on verbose
    start = time_idx[-1]
    forecast_idx = pd.date_range(start=start, periods=10, freq="YS-JAN")
    scenario = pd.DataFrame(0, index=forecast_idx, columns=[0, 1, 2])

    ss_mod._validate_forecast_args(
        time_index=time_idx, start=start, periods=10, scenario=scenario, use_scenario_index=True
    )
    last_message = caplog.messages[-1]
    assert "start, end, and periods arguments are ignored" in last_message

    # Verbose=False silences warning
    ss_mod._validate_forecast_args(
        time_index=time_idx,
        start=start,
        periods=10,
        scenario=scenario,
        use_scenario_index=True,
        verbose=False,
    )
    assert len(caplog.messages) == 1


@pytest.mark.parametrize("use_datetime_index", [True, False])
def test_forecast_index(use_datetime_index):
    ss_mod = make_statespace_mod(
        k_endog=1, k_posdef=1, k_states=2, filter_type="standard", verbose=False
    )
    ss_mod._fit_coords = dict()
    time_idx = _make_time_idx(ss_mod, use_datetime_index)

    # From start and end
    start = time_idx[-1]
    delta = pd.DateOffset(years=10) if use_datetime_index else 11
    end = start + delta

    x0_index, forecast_idx = ss_mod._build_forecast_index(time_idx, start=start, end=end)
    assert start not in forecast_idx
    assert x0_index == start
    assert forecast_idx.shape == (10,)

    # From start and periods
    start = time_idx[-1]
    periods = 10

    x0_index, forecast_idx = ss_mod._build_forecast_index(time_idx, start=start, periods=periods)
    assert start not in forecast_idx
    assert x0_index == start
    assert forecast_idx.shape == (10,)

    # From integer start
    start = 10
    x0_index, forecast_idx = ss_mod._build_forecast_index(time_idx, start=start, periods=periods)
    delta = forecast_idx.freq if use_datetime_index else 1

    assert x0_index == time_idx[start]
    assert forecast_idx.shape == (10,)
    assert (forecast_idx == time_idx[start + 1 : start + periods + 1]).all()

    # From scenario index
    scenario = pd.DataFrame(0, index=forecast_idx, columns=[0, 1, 2])
    new_start, forecast_idx = ss_mod._build_forecast_index(
        time_index=time_idx, scenario=scenario, use_scenario_index=True
    )
    assert x0_index not in forecast_idx
    assert x0_index == (forecast_idx[0] - delta)
    assert forecast_idx.shape == (10,)
    assert forecast_idx.equals(scenario.index)

    # From dictionary of scenarios
    scenario = {"a": pd.DataFrame(0, index=forecast_idx, columns=[0, 1, 2])}
    x0_index, forecast_idx = ss_mod._build_forecast_index(
        time_index=time_idx, scenario=scenario, use_scenario_index=True
    )
    assert x0_index == (forecast_idx[0] - delta)
    assert forecast_idx.shape == (10,)
    assert forecast_idx.equals(scenario["a"].index)


@pytest.mark.parametrize(
    "data_type",
    [pd.Series, pd.DataFrame, np.array, list, tuple],
    ids=["series", "dataframe", "array", "list", "tuple"],
)
def test_validate_scenario(data_type):
    if data_type is pd.DataFrame:
        # Ensure dataframes have the correct column name
        data_type = partial(pd.DataFrame, columns=["column_1"])

    # One data case
    data_info = {"a": {"shape": (None, 1), "dims": ("time", "features_a")}}
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(features_a=["column_1"])

    scenario = data_type(np.zeros(10))
    scenario = ss_mod._validate_scenario_data(scenario)

    # Lists and tuples are cast to 2d arrays
    if data_type in [tuple, list]:
        assert isinstance(scenario, np.ndarray)
        assert scenario.shape == (10, 1)

    # A one-item dictionary should also work
    scenario = {"a": scenario}
    ss_mod._validate_scenario_data(scenario)

    # Now data has to be a dictionary
    data_info.update({"b": {"shape": (None, 1), "dims": ("time", "features_b")}})
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(features_a=["column_1"], features_b=["column_1"])

    scenario = {"a": data_type(np.zeros(10)), "b": data_type(np.zeros(10))}
    ss_mod._validate_scenario_data(scenario)

    # Mixed data types
    data_info.update({"a": {"shape": (None, 10), "dims": ("time", "features_a")}})
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(
        features_a=[f"column_{i}" for i in range(10)], features_b=["column_1"]
    )

    scenario = {
        "a": pd.DataFrame(np.zeros((10, 10)), columns=ss_mod._fit_coords["features_a"]),
        "b": data_type(np.arange(10)),
    }

    ss_mod._validate_scenario_data(scenario)


@pytest.mark.parametrize(
    "data_type",
    [pd.Series, pd.DataFrame, np.array, list, tuple],
    ids=["series", "dataframe", "array", "list", "tuple"],
)
@pytest.mark.parametrize("use_datetime_index", [True, False])
def test_finalize_scenario_single(data_type, use_datetime_index):
    if data_type is pd.DataFrame:
        # Ensure dataframes have the correct column name
        data_type = partial(pd.DataFrame, columns=["column_1"])

    data_info = {"a": {"shape": (None, 1), "dims": ("time", "features_a")}}
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(features_a=["column_1"])

    time_idx = _make_time_idx(ss_mod, use_datetime_index)

    scenario = data_type(np.zeros((10,)))

    scenario = ss_mod._validate_scenario_data(scenario)
    t0, forecast_idx = ss_mod._build_forecast_index(time_idx, start=time_idx[-1], periods=10)
    scenario = ss_mod._finalize_scenario_initialization(scenario, forecast_index=forecast_idx)

    assert isinstance(scenario, pd.DataFrame)
    assert scenario.index.equals(forecast_idx)
    assert scenario.columns == ["column_1"]


@pytest.mark.parametrize(
    "data_type",
    [pd.Series, pd.DataFrame, np.array, list, tuple],
    ids=["series", "dataframe", "array", "list", "tuple"],
)
@pytest.mark.parametrize("use_datetime_index", [True, False])
@pytest.mark.parametrize("use_scenario_index", [True, False])
def test_finalize_secenario_dict(data_type, use_datetime_index, use_scenario_index):
    data_info = {
        "a": {"shape": (None, 1), "dims": ("time", "features_a")},
        "b": {"shape": (None, 2), "dims": ("time", "features_b")},
    }
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(features_a=["column_1"], features_b=["column_1", "column_2"])
    time_idx = _make_time_idx(ss_mod, use_datetime_index)

    initial_index = (
        pd.date_range(start=time_idx[-1], periods=10, freq=time_idx.freq)
        if use_datetime_index
        else pd.RangeIndex(time_idx[-1], time_idx[-1] + 10, 1)
    )

    if data_type is pd.DataFrame:
        # Ensure dataframes have the correct column name
        data_type = partial(pd.DataFrame, columns=["column_1"], index=initial_index)
    elif data_type is pd.Series:
        data_type = partial(pd.Series, index=initial_index)

    scenario = {
        "a": data_type(np.zeros((10,))),
        "b": pd.DataFrame(
            np.zeros((10, 2)), columns=ss_mod._fit_coords["features_b"], index=initial_index
        ),
    }

    scenario = ss_mod._validate_scenario_data(scenario)

    if use_scenario_index and data_type not in [np.array, list, tuple]:
        t0, forecast_idx = ss_mod._build_forecast_index(
            time_idx, scenario=scenario, periods=10, use_scenario_index=True
        )
    elif use_scenario_index and data_type in [np.array, list, tuple]:
        t0, forecast_idx = ss_mod._build_forecast_index(
            time_idx, scenario=scenario, start=-1, periods=10, use_scenario_index=True
        )
    else:
        t0, forecast_idx = ss_mod._build_forecast_index(time_idx, start=time_idx[-1], periods=10)

    scenario = ss_mod._finalize_scenario_initialization(scenario, forecast_index=forecast_idx)

    assert list(scenario.keys()) == ["a", "b"]
    assert all(isinstance(value, pd.DataFrame) for value in scenario.values())
    assert all(value.index.equals(forecast_idx) for value in scenario.values())


def test_invalid_scenarios():
    data_info = {"a": {"shape": (None, 1), "dims": ("time", "features_a")}}
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(features_a=["column_1", "column_2"])

    # Omitting the data raises
    with pytest.raises(
        ValueError, match="This model was fit using exogenous data. Forecasting cannot be performed"
    ):
        ss_mod._validate_scenario_data(None)

    # Giving a list, tuple, or Series when a matrix of data is expected should always raise
    with pytest.raises(
        ValueError,
        match="Scenario data for variable 'a' has the wrong number of columns. "
        "Expected 2, got 1",
    ):
        for data_type in [list, tuple, pd.Series]:
            ss_mod._validate_scenario_data(data_type(np.zeros(10)))
            ss_mod._validate_scenario_data({"a": data_type(np.zeros(10))})

    # Providing irrevelant data raises
    with pytest.raises(
        ValueError,
        match="Scenario data provided for variable 'jk lol', which is not an exogenous " "variable",
    ):
        ss_mod._validate_scenario_data({"jk lol": np.zeros(10)})

    # Incorrect 2nd dimension of a non-dataframe
    with pytest.raises(
        ValueError,
        match="Scenario data for variable 'a' has the wrong number of columns. Expected "
        "2, got 1",
    ):
        scenario = np.zeros(10).tolist()
        ss_mod._validate_scenario_data(scenario)
        ss_mod._validate_scenario_data(tuple(scenario))

        scenario = {"a": np.zeros(10).tolist()}
        ss_mod._validate_scenario_data(scenario)
        ss_mod._validate_scenario_data({"a": tuple(scenario["a"])})

    # If a data frame is provided, it needs to have all columns
    with pytest.raises(
        ValueError, match="Scenario data for variable 'a' is missing the following column: column_2"
    ):
        scenario = pd.DataFrame(np.zeros((10, 1)), columns=["column_1"])
        ss_mod._validate_scenario_data(scenario)

    # Extra columns also raises
    with pytest.raises(
        ValueError,
        match="Scenario data for variable 'a' contains the following extra columns "
        "that are not used by the model: column_3, column_4",
    ):
        scenario = pd.DataFrame(
            np.zeros((10, 4)), columns=["column_1", "column_2", "column_3", "column_4"]
        )
        ss_mod._validate_scenario_data(scenario)

    # Wrong number of time steps raises
    data_info = {
        "a": {"shape": (None, 1), "dims": ("time", "features_a")},
        "b": {"shape": (None, 1), "dims": ("time", "features_b")},
    }
    ss_mod = make_statespace_mod(
        k_endog=1,
        k_posdef=1,
        k_states=2,
        filter_type="standard",
        verbose=False,
        data_info=data_info,
    )
    ss_mod._fit_coords = dict(
        features_a=["column_1", "column_2"], features_b=["column_1", "column_2"]
    )

    with pytest.raises(
        ValueError, match="Scenario data must have the same number of time steps for all variables"
    ):
        scenario = {
            "a": pd.DataFrame(np.zeros((10, 2)), columns=ss_mod._fit_coords["features_a"]),
            "b": pd.DataFrame(np.zeros((11, 2)), columns=ss_mod._fit_coords["features_b"]),
        }
        ss_mod._validate_scenario_data(scenario)


@pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
@pytest.mark.parametrize("filter_output", ["predicted", "filtered", "smoothed"])
@pytest.mark.parametrize(
    "mod_name, idata_name, start, end, periods",
    [
        ("ss_mod_no_exog", "idata_no_exog", None, None, 10),
        ("ss_mod_no_exog", "idata_no_exog", -1, None, 10),
        ("ss_mod_no_exog", "idata_no_exog", 10, None, 10),
        ("ss_mod_no_exog", "idata_no_exog", 10, 21, None),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", None, None, 10),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", -1, None, 10),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", 10, None, 10),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", 10, "2020-01-21", None),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", "2020-03-01", "2020-03-11", None),
        ("ss_mod_no_exog_dt", "idata_no_exog_dt", "2020-03-01", None, 10),
    ],
    ids=[
        "range_default",
        "range_negative",
        "range_int",
        "range_end",
        "datetime_default",
        "datetime_negative",
        "datetime_int",
        "datetime_int_end",
        "datetime_datetime_end",
        "datetime_datetime",
    ],
)
def test_forecast(filter_output, mod_name, idata_name, start, end, periods, rng, request):
    mod = request.getfixturevalue(mod_name)
    idata = request.getfixturevalue(idata_name)
    time_idx = mod._get_fit_time_index()
    is_datetime = isinstance(time_idx, pd.DatetimeIndex)

    if isinstance(start, str):
        t0 = pd.Timestamp(start)
    elif isinstance(start, int):
        t0 = time_idx[start]
    else:
        t0 = time_idx[-1]

    delta = time_idx.freq if is_datetime else 1

    forecast_idata = mod.forecast(
        idata, start=start, end=end, periods=periods, filter_output=filter_output, random_seed=rng
    )

    forecast_idx = forecast_idata.coords["time"].values
    forecast_idx = pd.DatetimeIndex(forecast_idx) if is_datetime else pd.Index(forecast_idx)

    assert forecast_idx.shape == (10,)
    assert forecast_idata.forecast_latent.dims == ("chain", "draw", "time", "state")
    assert forecast_idata.forecast_observed.dims == ("chain", "draw", "time", "observed_state")

    assert not np.any(np.isnan(forecast_idata.forecast_latent.values))
    assert not np.any(np.isnan(forecast_idata.forecast_observed.values))

    assert forecast_idx[0] == (t0 + delta)


@pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
@pytest.mark.parametrize("start", [None, -1, 10])
def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
    scenario = pd.DataFrame(np.zeros((10, 3)), columns=["a", "b", "c"])
    scenario.iloc[5, 0] = 1e9

    forecast_idata = exog_ss_mod.forecast(
        idata_exog, start=start, periods=10, random_seed=rng, scenario=scenario
    )

    components = exog_ss_mod.extract_components_from_idata(forecast_idata)
    level = components.forecast_latent.sel(state="LevelTrend[level]")
    betas = components.forecast_latent.sel(state=["exog[a]", "exog[b]", "exog[c]"])

    scenario.index.name = "time"
    scenario_xr = (
        scenario.unstack()
        .to_xarray()
        .rename({"level_0": "state"})
        .assign_coords(state=["exog[a]", "exog[b]", "exog[c]"])
    )

    regression_effect = forecast_idata.forecast_observed.isel(observed_state=0) - level
    regression_effect_expected = (betas * scenario_xr).sum(dim=["state"])

    assert_allclose(regression_effect, regression_effect_expected)
