from collections.abc import Sequence
from typing import Any

import numpy as np
import pytensor.tensor as pt

from pytensor import graph_replace
from pytensor.tensor.slinalg import solve_discrete_lyapunov

from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
from pymc_extras.statespace.models.utilities import make_default_coords
from pymc_extras.statespace.utils.constants import (
    ALL_STATE_AUX_DIM,
    ALL_STATE_DIM,
    ETS_SEASONAL_DIM,
    OBS_STATE_AUX_DIM,
    OBS_STATE_DIM,
)


class BayesianETS(PyMCStateSpace):
    def __init__(
        self,
        order: tuple[str, str, str] | None = None,
        endog_names: str | list[str] | None = None,
        k_endog: int = 1,
        trend: bool = True,
        damped_trend: bool = False,
        seasonal: bool = False,
        seasonal_periods: int | None = None,
        measurement_error: bool = False,
        use_transformed_parameterization: bool = False,
        dense_innovation_covariance: bool = False,
        stationary_initialization: bool = False,
        initialization_dampening: float = 0.8,
        filter_type: str = "standard",
        verbose: bool = True,
    ):
        r"""
        Exponential Smoothing State Space Model

        This class can represent a subset of exponential smoothing state space models, specifically those with additive
        errors. Following .. [1], The general form of the model is:

        .. math::

            \begin{align}
            y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
            \epsilon_t &\sim N(0, \sigma)
            \end{align}

        where :math:`l_t` is the level component, :math:`b_t` is the trend component, and :math:`s_t` is the seasonal
        component. These components can be included or excluded, leading to different model specifications. The following
        models are possible:

        * `ETS(A,N,N)`: Simple exponential smoothing

            .. math::

                \begin{align}
                y_t &= l_{t-1} + \epsilon_t \\
                l_t &= l_{t-1} + \alpha \epsilon_t
                \end{align}

        Where :math:`\alpha \in [0, 1]` is a mixing parameter between past observations and current innovations.
        These equations arise by starting from the "component form":

            .. math::

                \begin{align}
                \hat{y}_{t+1 | t} &= l_t \\
                l_t &= \alpha y_t + (1 - \alpha) l_{t-1} \\
                &= l_{t-1} + \alpha (y_t - l_{t-1})
                &= l_{t-1} + \alpha \epsilon_t
                \end{align}

        Where $\epsilon_t$ are the forecast errors, assumed to be IID mean zero and normally distributed. The role of
        :math:`\alpha` is clearest in the second line. The level of the time series at each time is a mixture of
        :math:`\alpha` percent of the incoming data, and :math:`1 - \alpha` percent of the previous level. Recursive
        substitution reveals that the level is a weighted composite of all previous observations; thus the name
        "Exponential Smoothing".

        Additional supposed specifications include:

        * `ETS(A,A,N)`: Holt's linear trend method

            .. math::

                \begin{align}
                y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
                l_t &= l_{t-1} + b_{t-1} + \alpha \epsilon_t \\
                b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t
                \end{align}

            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.

        * `ETS(A,N,A)`: Additive seasonal method

            .. math::

                \begin{align}
                y_t &= l_{t-1} + s_{t-m} + \epsilon_t \\
                l_t &= l_{t-1} + \alpha \epsilon_t \\
                s_t &= s_{t-m} + (1 - \alpha)\gamma^\star \epsilon_t
                \end{align}

            [1]_ also consider an alternative parameterization with :math:`\gamma = (1 - \alpha) \gamma^\star`.

        * `ETS(A,A,A)`: Additive Holt-Winters method

            .. math::

                \begin{align}
                y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
                l_t &= l_{t-1} + \alpha \epsilon_t \\
                b_t &= b_{t-1} + \alpha \beta^\star \epsilon_t \\
                s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
                \end{align}

            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
            :math:`\gamma = (1 - \alpha) \gamma^\star`.

        * `ETS(A, Ad, N)`: Dampened trend method

            .. math::

                \begin{align}
                y_t &= l_{t-1} + b_{t-1} + \epsilon_t \\
                l_t &= l_{t-1} + \alpha \epsilon_t \\
                b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t
                \end{align}

            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^\star`.

        * `ETS(A, Ad, A)`: Dampened trend with seasonal method

            .. math::

                \begin{align}
                y_t &= l_{t-1} + b_{t-1} + s_{t-m} + \epsilon_t \\
                l_t &= l_{t-1} + \alpha \epsilon_t \\
                b_t &= \phi b_{t-1} + \alpha \beta^\star \epsilon_t \\
                s_t &= s_{t-m} + (1 - \alpha) \gamma^\star \epsilon_t
                \end{align}

            [1]_ also consider an alternative parameterization with :math:`\beta = \alpha \beta^star` and
            :math:`\gamma = (1 - \alpha) \gamma^\star`.


        Parameters
        ----------
        order: tuple of string, Optional
            The exponential smoothing "order". This is a tuple of three strings, each of which should be one of 'A', 'Ad',
            or 'N'.
            If provided, the model will be initialized from the given order, and the `trend`, `damped_trend`, and `seasonal`
            arguments will be ignored.
        endog_names: str or list of str, Optional
            Names associated with observed states. If a list, the length should be equal to the number of time series
            to be estimated.
        k_endog: int, Optional
            Number of time series to estimate. If endog_names are provided, this is ignored and len(endog_names) is
            used instead.
        trend: bool
            Whether to include a trend component. Setting ``trend=True`` is equivalent to ``order[1] == 'A'``.
        damped_trend: bool
            Whether to include a damping parameter on the trend component. Ignored if `trend` is `False`. Setting
            ``trend=True`` and ``damped_trend=True`` is equivalent to order[1] == 'Ad'.
        seasonal: bool
            Whether to include a seasonal component. Setting ``seasonal=True`` is equivalent to ``order[2] = 'A'``.
        seasonal_periods: int
            The number of periods in a complete seasonal cycle. Ignored if `seasonal` is `False`
            (or if ``order[2] == "N"``)
        measurement_error: bool
            Whether to include a measurement error term in the model. Default is `False`.
        use_transformed_parameterization: bool, default False
            If true, use the :math:`\alpha, \beta, \gamma` parameterization, otherwise use the :math:`\alpha, \beta^\star,
            \gamma^\star` parameterization. This will change the admissible region for the priors.

            - Under the **non-transformed** parameterization, all of :math:`\alpha, \beta^\star, \gamma^\star` should be
              between 0 and 1.
            - Under the **transformed**  parameterization, :math:`\alpha \in (0, 1)`, :math:`\beta \in (0, \alpha)`, and
              :math:`\gamma \in (0, 1 - \alpha)`

            The :meth:`param_info` method will change to reflect the suggested intervals based on the value of this
            argument.
        dense_innovation_covariance: bool, default False
            Whether to estimate a dense covariance for statespace innovations. In an ETS models, each observed variable
            has a single source of stochastic variation. If True, these innovations are allowed to be correlated.
            Ignored if ``k_endog == 1``
        stationary_initialization: bool, default False
            If True, the Kalman Filter's initial covariance matrix will be set to an approximate steady-state value.
            The approximation is formed by adding a small dampening factor to each state. Specifically, the level state
            for a ('A', 'N', 'N') model is written:

            .. math::
                \ell_t = \ell_{t-1} + \alpha * e_t

            That this system is not stationary can be understood in ARIMA terms: the level is a random walk; that is,
            :math:`rho = 1`. This can be remedied by pretending that we instead have a dampened system:

            .. math::
                \ell_t = \rho \ell_{t-1} + \alpha * e_t

            With :math:`\rho \approx 1`, the system is stationary, and we can solve for the steady-state covariance
            matrix. This is then used as the initial covariance matrix for the Kalman Filter. This is a heuristic
            method that helps avoid setting a prior on the initial covariance matrix.
        initialization_dampening: float, default 0.8
            Dampening factor to add to non-stationary model components. This is only used for initialization, it does
            *not* add dampening to the model. Ignored if `stationary_initialization` is `False`.
        filter_type: str, default "standard"
            The type of Kalman Filter to use. Options are "standard", "single", "univariate", "steady_state",
            and "cholesky". See the docs for kalman filters for more details.
        verbose: bool, default True
            If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.

        References
        ----------
        .. [1] Hyndman, Rob J., and George Athanasopoulos. Forecasting: principles and practice. OTexts, 2018.
        """

        if order is not None:
            if len(order) != 3 or any(not isinstance(o, str) for o in order):
                raise ValueError("Order must be a tuple of three strings.")
            if order[0] != "A":
                raise ValueError("Only additive errors are supported.")
            if order[1] not in {"A", "Ad", "N"}:
                raise ValueError(
                    f"Invalid trend specification. Only 'A' (additive), 'Ad' (additive with dampening), "
                    f"or 'N' (no trend) are allowed. Found {order[1]}"
                )
            if order[2] not in {"A", "N"}:
                raise ValueError(
                    f"Invalid seasonal specification. Only 'A' (additive) or 'N' (no seasonal component) "
                    f"are allowed. Found {order[2]}"
                )

            trend = order[1] != "N"
            damped_trend = order[1] == "Ad"
            seasonal = order[2] == "A"

        self.trend = trend
        self.damped_trend = damped_trend
        self.seasonal = seasonal
        self.seasonal_periods = seasonal_periods
        self.use_transformed_parameterization = use_transformed_parameterization
        self.stationary_initialization = stationary_initialization

        if not (0.0 < initialization_dampening < 1.0):
            raise ValueError(
                "Dampening term used for initialization must be between 0 and 1 (preferably close to"
                "1.0)"
            )

        self.initialization_dampening = initialization_dampening

        if self.seasonal and self.seasonal_periods is None:
            raise ValueError("If seasonal is True, seasonal_periods must be provided.")

        if endog_names is not None:
            endog_names = list(endog_names)
            k_endog = len(endog_names)
        else:
            endog_names = [f"data_{i}" for i in range(k_endog)] if k_endog > 1 else ["data"]

        self.endog_names = endog_names

        if dense_innovation_covariance and k_endog == 1:
            dense_innovation_covariance = False

        self.dense_innovation_covariance = dense_innovation_covariance

        k_states = (
            2
            + int(trend)
            + int(seasonal) * (seasonal_periods if seasonal_periods is not None else 0)
        ) * k_endog

        k_posdef = k_endog

        super().__init__(
            k_endog,
            k_states,
            k_posdef,
            filter_type,
            verbose=verbose,
            measurement_error=measurement_error,
        )

    @property
    def param_names(self):
        names = [
            "initial_level",
            "initial_trend",
            "initial_seasonal",
            "P0",
            "alpha",
            "beta",
            "gamma",
            "phi",
            "sigma_state",
            "state_cov",
            "sigma_obs",
        ]
        if not self.trend:
            names.remove("initial_trend")
            names.remove("beta")
        if not self.damped_trend:
            names.remove("phi")
        if not self.seasonal:
            names.remove("initial_seasonal")
            names.remove("gamma")
        if not self.measurement_error:
            names.remove("sigma_obs")

        if self.dense_innovation_covariance:
            names.remove("sigma_state")
        else:
            names.remove("state_cov")

        if self.stationary_initialization:
            names.remove("P0")

        return names

    @property
    def param_info(self) -> dict[str, dict[str, Any]]:
        info = {
            "P0": {
                "shape": (self.k_states, self.k_states),
                "constraints": "Positive Semi-definite",
            },
            "initial_level": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": None,
            },
            "initial_trend": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": None,
            },
            "initial_seasonal": {"shape": (self.seasonal_periods,), "constraints": None},
            "sigma_obs": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": "Positive",
            },
            "sigma_state": {
                "shape": None if self.k_posdef == 1 else (self.k_posdef,),
                "constraints": "Positive",
            },
            "alpha": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": "0 < alpha < 1",
            },
            "beta": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": "0 < beta < 1"
                if not self.use_transformed_parameterization
                else "0 < beta < alpha",
            },
            "gamma": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": "0 < gamma< 1"
                if not self.use_transformed_parameterization
                else "0 < gamma < (1 - alpha)",
            },
            "phi": {
                "shape": None if self.k_endog == 1 else (self.k_endog,),
                "constraints": "0 < phi < 1",
            },
        }

        if self.dense_innovation_covariance:
            del info["sigma_state"]
            info["state_cov"] = {
                "shape": (self.k_posdef, self.k_posdef),
                "constraints": "Positive Semi-definite",
            }

        for name in self.param_names:
            info[name]["dims"] = self.param_dims.get(name, None)

        return {name: info[name] for name in self.param_names}

    @property
    def state_names(self):
        states = ["innovation", "level"]
        if self.trend:
            states += ["trend"]
        if self.seasonal:
            states += ["seasonality"]
            states += [f"L{i}.season" for i in range(1, self.seasonal_periods)]

        if self.k_endog > 1:
            states = [f"{name}_{state}" for name in self.endog_names for state in states]

        return states

    @property
    def observed_states(self):
        return self.endog_names

    @property
    def shock_names(self):
        return (
            ["innovation"]
            if self.k_endog == 1
            else [f"{name}_innovation" for name in self.endog_names]
        )

    @property
    def param_dims(self):
        coord_map = {
            "P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
            "sigma_obs": (OBS_STATE_DIM,),
            "sigma_state": (OBS_STATE_DIM,),
            "initial_level": (OBS_STATE_DIM,),
            "initial_trend": (OBS_STATE_DIM,),
            "initial_seasonal": (ETS_SEASONAL_DIM,),
            "seasonal_param": (ETS_SEASONAL_DIM,),
        }

        if self.dense_innovation_covariance:
            del coord_map["sigma_state"]
            coord_map["state_cov"] = (OBS_STATE_DIM, OBS_STATE_AUX_DIM)

        if self.k_endog == 1:
            coord_map["sigma_state"] = None
            coord_map["sigma_obs"] = None
            coord_map["initial_level"] = None
            coord_map["initial_trend"] = None
        else:
            coord_map["alpha"] = (OBS_STATE_DIM,)
            coord_map["beta"] = (OBS_STATE_DIM,)
            coord_map["gamma"] = (OBS_STATE_DIM,)
            coord_map["phi"] = (OBS_STATE_DIM,)
            coord_map["initial_seasonal"] = (OBS_STATE_DIM, ETS_SEASONAL_DIM)
            coord_map["seasonal_param"] = (OBS_STATE_DIM, ETS_SEASONAL_DIM)

        if not self.measurement_error:
            del coord_map["sigma_obs"]
        if not self.seasonal:
            del coord_map["seasonal_param"]

        return coord_map

    @property
    def coords(self) -> dict[str, Sequence]:
        coords = make_default_coords(self)
        if self.seasonal:
            coords.update({ETS_SEASONAL_DIM: list(range(1, self.seasonal_periods + 1))})

        return coords

    def _stationary_initialization(self, T_stationary):
        # Solve for matrix quadratic for P0
        R = self.ssm["selection"]
        Q = self.ssm["state_cov"]

        # ETS models are not stationary, but we can proceed *as if* the model were stationary by introducing large
        # dampening factors on all components. We then set the initial covariance to the steady-state of that system,
        # which we hope is similar enough to give a good initialization for the non-stationary system.

        T_stationary = pt.specify_shape(T_stationary, (self.k_states, self.k_states))
        P0 = solve_discrete_lyapunov(T_stationary, pt.linalg.matrix_dot(R, Q, R.T))
        P0 = pt.specify_shape(P0, (self.k_states, self.k_states))

        return P0

    def make_symbolic_graph(self) -> None:
        k_states_each = self.k_states // self.k_endog

        initial_level = self.make_and_register_variable(
            "initial_level", shape=(self.k_endog,) if self.k_endog > 1 else (), dtype=floatX
        )

        initial_states = [pt.zeros(k_states_each) for _ in range(self.k_endog)]
        if self.k_endog == 1:
            initial_states = [pt.set_subtensor(initial_states[0][1], initial_level)]
        else:
            initial_states = [
                pt.set_subtensor(initial_state[1], initial_level[i])
                for i, initial_state in enumerate(initial_states)
            ]

        # The shape of R can be pre-allocated, then filled with the required parameters
        R = pt.zeros((self.k_states // self.k_endog, 1))

        alpha = self.make_and_register_variable(
            "alpha", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
        )

        # This is a dummy value for initialization. When we do a stationary initialization, it will be set to a value
        # close to 1. Otherwise, it will be 1. We do not want this value to exist outside of this method.
        stationary_dampening = pt.scalar("dampen_dummy")

        if self.k_endog == 1:
            # The R[0, 0] entry needs to be adjusted for a shift in the time indices. Consider the (A, N, N) model:
            # y_t = l_{t-1} + e_t
            # l_t = l_{t-1} + alpha * e_t
            R_list = [pt.set_subtensor(R[1, 0], alpha)]  # and l_t = ... + alpha * e_t

            # We want the first equation to be in terms of time t on the RHS, because our observation equation is always
            # y_t = Z @ x_t. Re-arranging equation 2, we get l_{t-1} = l_t - alpha * e_t --> y_t = l_t + e_t - alpha * e_t
            # --> y_t = l_t + (1 - alpha) * e_t
            R_list = [pt.set_subtensor(R[0, :], (1 - alpha)) for R in R_list]
        else:
            # If there are multiple endog, clone the basic R matrix and modify the appropriate entries
            R_list = [pt.set_subtensor(R[1, 0], alpha[i]) for i in range(self.k_endog)]
            R_list = [pt.set_subtensor(R[0, :], (1 - alpha[i])) for i, R in enumerate(R_list)]

        # Shock and level component always exists, the base case is e_t = e_t and l_t = l_{t-1}
        T_base = pt.set_subtensor(pt.zeros((2, 2))[1, 1], stationary_dampening)

        if self.trend:
            initial_trend = self.make_and_register_variable(
                "initial_trend", shape=(self.k_endog,) if self.k_endog > 1 else (), dtype=floatX
            )

            if self.k_endog == 1:
                initial_states = [pt.set_subtensor(initial_states[0][2], initial_trend)]
            else:
                initial_states = [
                    pt.set_subtensor(initial_state[2], initial_trend[i])
                    for i, initial_state in enumerate(initial_states)
                ]
            beta = self.make_and_register_variable(
                "beta", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
            )
            if self.use_transformed_parameterization:
                param = beta
            else:
                param = alpha * beta
            if self.k_endog == 1:
                R_list = [pt.set_subtensor(R[2, 0], param) for R in R_list]
            else:
                R_list = [pt.set_subtensor(R[2, 0], param[i]) for i, R in enumerate(R_list)]

            # If a trend is requested, we have the following transition equations (omitting the shocks):
            # l_t = l_{t-1} + b_{t-1}
            # b_t = b_{t-1}
            T_base = pt.as_tensor_variable(([0.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]))
            T_base = pt.set_subtensor(T_base[[1, 2], [1, 2]], stationary_dampening)

        if self.damped_trend:
            phi = self.make_and_register_variable(
                "phi", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
            )
            # We are always in the case where we have a trend, so we can add the dampening parameter to T_base defined
            # in that branch. Transition equations become:
            # l_t = l_{t-1} + phi * b_{t-1}
            # b_t = phi * b_{t-1}
            if self.k_endog > 1:
                T_base = [pt.set_subtensor(T_base[1:, 2], phi[i]) for i in range(self.k_endog)]
            else:
                T_base = pt.set_subtensor(T_base[1:, 2], phi)

        T_components = (
            [T_base for _ in range(self.k_endog)] if not isinstance(T_base, list) else T_base
        )

        if self.seasonal:
            initial_seasonal = self.make_and_register_variable(
                "initial_seasonal",
                shape=(self.seasonal_periods,)
                if self.k_endog == 1
                else (self.k_endog, self.seasonal_periods),
                dtype=floatX,
            )
            if self.k_endog == 1:
                initial_states = [
                    pt.set_subtensor(initial_states[0][2 + int(self.trend) :], initial_seasonal)
                ]
            else:
                initial_states = [
                    pt.set_subtensor(initial_state[2 + int(self.trend) :], initial_seasonal[i])
                    for i, initial_state in enumerate(initial_states)
                ]

            gamma = self.make_and_register_variable(
                "gamma", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
            )

            param = gamma if self.use_transformed_parameterization else (1 - alpha) * gamma
            # Additional adjustment to the R[0, 0] position is required. Start from:
            # y_t = l_{t-1} + s_{t-m} + e_t
            # l_t = l_{t-1} + alpha * e_t
            # s_t = s_{t-m} + gamma * e_t
            # Solve for l_{t-1} and s_{t-m} in terms of l_t and s_t, then substitute into the observation equation:
            # y_t = l_t + s_t - alpha * e_t - gamma * e_t + e_t --> y_t = l_t + s_t + (1 - alpha - gamma) * e_t

            if self.k_endog == 1:
                R_list = [pt.set_subtensor(R[2 + int(self.trend), 0], param) for R in R_list]
                R_list = [pt.set_subtensor(R[0, 0], R[0, 0] - param) for R in R_list]

            else:
                R_list = [
                    pt.set_subtensor(R[2 + int(self.trend), 0], param[i])
                    for i, R in enumerate(R_list)
                ]
                R_list = [
                    pt.set_subtensor(R[0, 0], R[0, 0] - param[i]) for i, R in enumerate(R_list)
                ]

            # The seasonal component is always going to look like a TimeFrequency structural component, see that
            # docstring for more details
            T_seasonals = [pt.eye(self.seasonal_periods, k=-1) for _ in range(self.k_endog)]
            T_seasonals = [
                pt.set_subtensor(T_seasonal[0, -1], stationary_dampening)
                for T_seasonal in T_seasonals
            ]

            # Organize the components so it goes T1, T_seasonal_1, T2, T_seasonal_2, etc.
            T_components = [
                matrix[i] for i in range(self.k_endog) for matrix in [T_components, T_seasonals]
            ]

        x0 = pt.concatenate(initial_states, axis=0)
        R = pt.linalg.block_diag(*R_list)

        self.ssm["initial_state"] = x0
        self.ssm["selection"] = pt.specify_shape(R, shape=(self.k_states, self.k_posdef))

        T = pt.linalg.block_diag(*T_components)

        # Remove the stationary_dampening dummies before saving the transition matrix
        self.ssm["transition"] = pt.specify_shape(
            graph_replace(T, {stationary_dampening: 1.0}), (self.k_states, self.k_states)
        )

        Zs = [np.zeros((self.k_endog, self.k_states // self.k_endog)) for _ in range(self.k_endog)]
        for i, Z in enumerate(Zs):
            Z[i, 0] = 1.0  # innovation
            Z[i, 1] = 1.0  # level
            if self.seasonal:
                Z[i, 2 + int(self.trend)] = 1.0

        Z = pt.concatenate(Zs, axis=1)

        self.ssm["design"] = Z

        # Set up the state covariance matrix
        if self.dense_innovation_covariance:
            state_cov = self.make_and_register_variable(
                "state_cov", shape=(self.k_posdef, self.k_posdef), dtype=floatX
            )
            self.ssm["state_cov"] = state_cov

        else:
            state_cov_idx = ("state_cov", *np.diag_indices(self.k_posdef))
            state_cov = self.make_and_register_variable(
                "sigma_state", shape=() if self.k_posdef == 1 else (self.k_posdef,), dtype=floatX
            )
            self.ssm[state_cov_idx] = state_cov**2

        if self.measurement_error:
            obs_cov_idx = ("obs_cov", *np.diag_indices(self.k_endog))
            obs_cov = self.make_and_register_variable(
                "sigma_obs", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
            )
            self.ssm[obs_cov_idx] = obs_cov**2

        if self.stationary_initialization:
            T_stationary = graph_replace(T, {stationary_dampening: self.initialization_dampening})
            P0 = self._stationary_initialization(T_stationary)

        else:
            P0 = self.make_and_register_variable(
                "P0", shape=(self.k_states, self.k_states), dtype=floatX
            )

        self.ssm["initial_state_cov"] = P0
