"""
Markov switching models

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

import warnings

import numpy as np
import pandas as pd
from scipy.special import logsumexp

from statsmodels.base.data import PandasData
import statsmodels.base.wrapper as wrap
from statsmodels.tools.decorators import cache_readonly
from statsmodels.tools.eval_measures import aic, bic, hqic
from statsmodels.tools.numdiff import approx_fprime_cs, approx_hess_cs
from statsmodels.tools.sm_exceptions import EstimationWarning
from statsmodels.tools.tools import Bunch, pinv_extended
import statsmodels.tsa.base.tsa_model as tsbase
from statsmodels.tsa.regime_switching._hamilton_filter import (
    chamilton_filter_log,
    dhamilton_filter_log,
    shamilton_filter_log,
    zhamilton_filter_log,
)
from statsmodels.tsa.regime_switching._kim_smoother import (
    ckim_smoother_log,
    dkim_smoother_log,
    skim_smoother_log,
    zkim_smoother_log,
)
from statsmodels.tsa.statespace.tools import (
    _safe_cond,
    find_best_blas_type,
    prepare_exog,
)

prefix_hamilton_filter_log_map = {
    "s": shamilton_filter_log,
    "d": dhamilton_filter_log,
    "c": chamilton_filter_log,
    "z": zhamilton_filter_log,
}

prefix_kim_smoother_log_map = {
    "s": skim_smoother_log,
    "d": dkim_smoother_log,
    "c": ckim_smoother_log,
    "z": zkim_smoother_log,
}


def _logistic(x):
    """
    Note that this is not a vectorized function
    """
    x = np.array(x)
    # np.exp(x) / (1 + np.exp(x))
    if x.ndim == 0:
        y = np.reshape(x, (1, 1, 1))
    # np.exp(x[i]) / (1 + np.sum(np.exp(x[:])))
    elif x.ndim == 1:
        y = np.reshape(x, (len(x), 1, 1))
    # np.exp(x[i,t]) / (1 + np.sum(np.exp(x[:,t])))
    elif x.ndim == 2:
        y = np.reshape(x, (x.shape[0], 1, x.shape[1]))
    # np.exp(x[i,j,t]) / (1 + np.sum(np.exp(x[:,j,t])))
    elif x.ndim == 3:
        y = x
    else:
        raise NotImplementedError

    tmp = np.c_[np.zeros((y.shape[-1], y.shape[1], 1)), y.T].T
    evaluated = np.reshape(np.exp(y - logsumexp(tmp, axis=0)), x.shape)

    return evaluated


def _partials_logistic(x):
    """
    Note that this is not a vectorized function
    """
    tmp = _logistic(x)

    # k
    if tmp.ndim == 0:
        return tmp - tmp**2
    # k x k
    elif tmp.ndim == 1:
        partials = np.diag(tmp - tmp**2)
    # k x k x t
    elif tmp.ndim == 2:
        partials = [np.diag(tmp[:, t] - tmp[:, t] ** 2) for t in range(tmp.shape[1])]
        shape = tmp.shape[1], tmp.shape[0], tmp.shape[0]
        partials = np.concatenate(partials).reshape(shape).transpose((1, 2, 0))
    # k x k x j x t
    else:
        partials = [
            [np.diag(tmp[:, j, t] - tmp[:, j, t] ** 2) for t in range(tmp.shape[2])]
            for j in range(tmp.shape[1])
        ]
        shape = tmp.shape[1], tmp.shape[2], tmp.shape[0], tmp.shape[0]
        partials = np.concatenate(partials).reshape(shape).transpose((2, 3, 0, 1))

    for i in range(tmp.shape[0]):
        for j in range(i):
            partials[i, j, ...] = -tmp[i, ...] * tmp[j, ...]
            partials[j, i, ...] = partials[i, j, ...]
    return partials


def cy_hamilton_filter_log(
    initial_probabilities, regime_transition, conditional_loglikelihoods, model_order
):
    """
    Hamilton filter in log space using Cython inner loop.

    Parameters
    ----------
    initial_probabilities : ndarray
        Array of initial probabilities, shaped (k_regimes,) giving the
        distribution of the regime process at time t = -order where order
        is a nonnegative integer.
    regime_transition : ndarray
        Matrix of regime transition probabilities, shaped either
        (k_regimes, k_regimes, 1) or if there are time-varying transition
        probabilities (k_regimes, k_regimes, nobs + order).  Entry [i, j,
        t] contains the probability of moving from j at time t-1 to i at
        time t, so each matrix regime_transition[:, :, t] should be left
        stochastic.  The first order entries and initial_probabilities are
        used to produce the initial joint distribution of dimension order +
        1 at time t=0.
    conditional_loglikelihoods : ndarray
        Array of loglikelihoods conditional on the last `order+1` regimes,
        shaped (k_regimes,)*(order + 1) + (nobs,).

    Returns
    -------
    filtered_marginal_probabilities : ndarray
        Array containing Pr[S_t=s_t | Y_t] - the probability of being in each
        regime conditional on time t information. Shaped (k_regimes, nobs).
    predicted_joint_probabilities : ndarray
        Array containing Pr[S_t=s_t, ..., S_{t-order}=s_{t-order} | Y_{t-1}] -
        the joint probability of the current and previous `order` periods
        being in each combination of regimes conditional on time t-1
        information. Shaped (k_regimes,) * (order + 1) + (nobs,).
    joint_loglikelihoods : ndarray
        Array of loglikelihoods condition on time t information,
        shaped (nobs,).
    filtered_joint_probabilities : ndarray
        Array containing Pr[S_t=s_t, ..., S_{t-order}=s_{t-order} | Y_{t}] -
        the joint probability of the current and previous `order` periods
        being in each combination of regimes conditional on time t
        information. Shaped (k_regimes,) * (order + 1) + (nobs,).
    """

    # Dimensions
    k_regimes = len(initial_probabilities)
    nobs = conditional_loglikelihoods.shape[-1]
    order = conditional_loglikelihoods.ndim - 2
    dtype = conditional_loglikelihoods.dtype

    # Check for compatible shapes.
    incompatible_shapes = (
        regime_transition.shape[-1] not in (1, nobs + model_order)
        or regime_transition.shape[:2] != (k_regimes, k_regimes)
        or conditional_loglikelihoods.shape[0] != k_regimes
    )
    if incompatible_shapes:
        raise ValueError("Arguments do not have compatible shapes")

    # Convert to log space
    initial_probabilities = np.log(initial_probabilities)
    regime_transition = np.log(np.maximum(regime_transition, 1e-20))

    # Storage
    # Pr[S_t = s_t | Y_t]
    filtered_marginal_probabilities = np.zeros((k_regimes, nobs), dtype=dtype)
    # Pr[S_t = s_t, ... S_{t-r} = s_{t-r} | Y_{t-1}]
    # Has k_regimes^(order+1) elements
    predicted_joint_probabilities = np.zeros(
        (k_regimes,) * (order + 1) + (nobs,), dtype=dtype
    )
    # log(f(y_t | Y_{t-1}))
    joint_loglikelihoods = np.zeros((nobs,), dtype)
    # Pr[S_t = s_t, ... S_{t-r+1} = s_{t-r+1} | Y_t]
    # Has k_regimes^order elements
    filtered_joint_probabilities = np.zeros(
        (k_regimes,) * (order + 1) + (nobs + 1,), dtype=dtype
    )

    # Initial probabilities
    filtered_marginal_probabilities[:, 0] = initial_probabilities
    tmp = np.copy(initial_probabilities)
    shape = (k_regimes, k_regimes)
    transition_t = 0
    for i in range(order):
        if regime_transition.shape[-1] > 1:
            transition_t = i
        tmp = np.reshape(regime_transition[..., transition_t], shape + (1,) * i) + tmp
    filtered_joint_probabilities[..., 0] = tmp

    # Get appropriate subset of transition matrix
    if regime_transition.shape[-1] > 1:
        regime_transition = regime_transition[..., model_order:]

    # Run Cython filter iterations
    prefix, dtype, _ = find_best_blas_type(
        (
            regime_transition,
            conditional_loglikelihoods,
            joint_loglikelihoods,
            predicted_joint_probabilities,
            filtered_joint_probabilities,
        )
    )
    func = prefix_hamilton_filter_log_map[prefix]
    func(
        nobs,
        k_regimes,
        order,
        regime_transition,
        conditional_loglikelihoods.reshape(k_regimes ** (order + 1), nobs),
        joint_loglikelihoods,
        predicted_joint_probabilities.reshape(k_regimes ** (order + 1), nobs),
        filtered_joint_probabilities.reshape(k_regimes ** (order + 1), nobs + 1),
    )

    # Save log versions for smoother
    predicted_joint_probabilities_log = predicted_joint_probabilities
    filtered_joint_probabilities_log = filtered_joint_probabilities

    # Convert out of log scale
    predicted_joint_probabilities = np.exp(predicted_joint_probabilities)
    filtered_joint_probabilities = np.exp(filtered_joint_probabilities)

    # S_t | t
    filtered_marginal_probabilities = filtered_joint_probabilities[..., 1:]
    for _ in range(1, filtered_marginal_probabilities.ndim - 1):
        filtered_marginal_probabilities = np.sum(
            filtered_marginal_probabilities, axis=-2
        )

    return (
        filtered_marginal_probabilities,
        predicted_joint_probabilities,
        joint_loglikelihoods,
        filtered_joint_probabilities[..., 1:],
        predicted_joint_probabilities_log,
        filtered_joint_probabilities_log[..., 1:],
    )


def cy_kim_smoother_log(
    regime_transition, predicted_joint_probabilities, filtered_joint_probabilities
):
    """
    Kim smoother in log space using Cython inner loop.

    Parameters
    ----------
    regime_transition : ndarray
        Matrix of regime transition probabilities, shaped either
        (k_regimes, k_regimes, 1) or if there are time-varying transition
        probabilities (k_regimes, k_regimes, nobs).
    predicted_joint_probabilities : ndarray
        Array containing Pr[S_t=s_t, ..., S_{t-order}=s_{t-order} | Y_{t-1}] -
        the joint probability of the current and previous `order` periods
        being in each combination of regimes conditional on time t-1
        information. Shaped (k_regimes,) * (order + 1) + (nobs,).
    filtered_joint_probabilities : ndarray
        Array containing Pr[S_t=s_t, ..., S_{t-order}=s_{t-order} | Y_{t}] -
        the joint probability of the current and previous `order` periods
        being in each combination of regimes conditional on time t
        information. Shaped (k_regimes,) * (order + 1) + (nobs,).

    Returns
    -------
    smoothed_joint_probabilities : ndarray
        Array containing Pr[S_t=s_t, ..., S_{t-order}=s_{t-order} | Y_T] -
        the joint probability of the current and previous `order` periods
        being in each combination of regimes conditional on all information.
        Shaped (k_regimes,) * (order + 1) + (nobs,).
    smoothed_marginal_probabilities : ndarray
        Array containing Pr[S_t=s_t | Y_T] - the probability of being in each
        regime conditional on all information. Shaped (k_regimes, nobs).
    """

    # Dimensions
    k_regimes = filtered_joint_probabilities.shape[0]
    nobs = filtered_joint_probabilities.shape[-1]
    order = filtered_joint_probabilities.ndim - 2
    dtype = filtered_joint_probabilities.dtype

    # Storage
    smoothed_joint_probabilities = np.zeros(
        (k_regimes,) * (order + 1) + (nobs,), dtype=dtype
    )

    # Get appropriate subset of transition matrix
    if regime_transition.shape[-1] == nobs + order:
        regime_transition = regime_transition[..., order:]

    # Convert to log space
    regime_transition = np.log(np.maximum(regime_transition, 1e-20))

    # Run Cython smoother iterations
    prefix, dtype, _ = find_best_blas_type(
        (regime_transition, predicted_joint_probabilities, filtered_joint_probabilities)
    )
    func = prefix_kim_smoother_log_map[prefix]
    func(
        nobs,
        k_regimes,
        order,
        regime_transition,
        predicted_joint_probabilities.reshape(k_regimes ** (order + 1), nobs),
        filtered_joint_probabilities.reshape(k_regimes ** (order + 1), nobs),
        smoothed_joint_probabilities.reshape(k_regimes ** (order + 1), nobs),
    )

    # Convert back from log space
    smoothed_joint_probabilities = np.exp(smoothed_joint_probabilities)

    # Get smoothed marginal probabilities S_t | T by integrating out
    # S_{t-k+1}, S_{t-k+2}, ..., S_{t-1}
    smoothed_marginal_probabilities = smoothed_joint_probabilities
    for _ in range(1, smoothed_marginal_probabilities.ndim - 1):
        smoothed_marginal_probabilities = np.sum(
            smoothed_marginal_probabilities, axis=-2
        )

    return smoothed_joint_probabilities, smoothed_marginal_probabilities


class MarkovSwitchingParams:
    """
    Class to hold parameters in Markov switching models

    Parameters
    ----------
    k_regimes : int
        The number of regimes between which parameters may switch.

    Notes
    -----

    The purpose is to allow selecting parameter indexes / slices based on
    parameter type, regime number, or both.

    Parameters are lexicographically ordered in the following way:

    1. Named type string (e.g. "autoregressive")
    2. Number (e.g. the first autoregressive parameter, then the second)
    3. Regime (if applicable)

    Parameter blocks are set using dictionary setter notation where the key
    is the named type string and the value is a list of boolean values
    indicating whether a given parameter is switching or not.

    For example, consider the following code:

        parameters = MarkovSwitchingParams(k_regimes=2)
        parameters['regime_transition'] = [1,1]
        parameters['exog'] = [0, 1]

    This implies the model has 7 parameters: 4 "regime_transition"-related
    parameters (2 parameters that each switch according to regimes) and 3
    "exog"-related parameters (1 parameter that does not switch, and one 1 that
    does).

    The order of parameters is then:

    1. The first "regime_transition" parameter, regime 0
    2. The first "regime_transition" parameter, regime 1
    3. The second "regime_transition" parameter, regime 1
    4. The second "regime_transition" parameter, regime 1
    5. The first "exog" parameter
    6. The second "exog" parameter, regime 0
    7. The second "exog" parameter, regime 1

    Retrieving indexes / slices is done through dictionary getter notation.
    There are three options for the dictionary key:

    - Regime number (zero-indexed)
    - Named type string (e.g. "autoregressive")
    - Regime number and named type string

    In the above example, consider the following getters:

    >>> parameters[0]
    array([0, 2, 4, 6])
    >>> parameters[1]
    array([1, 3, 5, 6])
    >>> parameters['exog']
    slice(4, 7, None)
    >>> parameters[0, 'exog']
    [4, 6]
    >>> parameters[1, 'exog']
    [4, 7]

    Notice that in the last two examples, both lists of indexes include 4.
    That's because that is the index of the the non-switching first "exog"
    parameter, which should be selected regardless of the regime.

    In addition to the getter, the `k_parameters` attribute is an dict
    with the named type strings as the keys. It can be used to get the total
    number of parameters of each type:

    >>> parameters.k_parameters['regime_transition']
    4
    >>> parameters.k_parameters['exog']
    3
    """

    def __init__(self, k_regimes):
        self.k_regimes = k_regimes

        self.k_params = 0
        self.k_parameters = {}
        self.switching = {}
        self.slices_purpose = {}
        self.relative_index_regime_purpose = [{} for i in range(self.k_regimes)]
        self.index_regime_purpose = [{} for i in range(self.k_regimes)]
        self.index_regime = [[] for i in range(self.k_regimes)]

    def __getitem__(self, key):
        _type = type(key)

        # Get a slice for a block of parameters by purpose
        if _type is str:
            return self.slices_purpose[key]
        # Get a slice for a block of parameters by regime
        elif _type is int:
            return self.index_regime[key]
        elif _type is tuple:
            if not len(key) == 2:
                raise IndexError("Invalid index")
            if type(key[1]) is str and type(key[0]) is int:
                return self.index_regime_purpose[key[0]][key[1]]
            elif type(key[0]) is str and type(key[1]) is int:
                return self.index_regime_purpose[key[1]][key[0]]
            else:
                raise IndexError("Invalid index")
        else:
            raise IndexError("Invalid index")

    def __setitem__(self, key, value):
        _type = type(key)

        if _type is str:
            value = np.array(value, dtype=bool, ndmin=1)
            k_params = self.k_params
            self.k_parameters[key] = value.size + np.sum(value) * (self.k_regimes - 1)
            self.k_params += self.k_parameters[key]
            self.switching[key] = value
            self.slices_purpose[key] = np.s_[k_params : self.k_params]

            for j in range(self.k_regimes):
                self.relative_index_regime_purpose[j][key] = []
                self.index_regime_purpose[j][key] = []

            offset = 0
            for i in range(value.size):
                switching = value[i]
                for j in range(self.k_regimes):
                    # Non-switching parameters
                    if not switching:
                        self.relative_index_regime_purpose[j][key].append(offset)
                    # Switching parameters
                    else:
                        self.relative_index_regime_purpose[j][key].append(offset + j)
                offset += 1 if not switching else self.k_regimes

            for j in range(self.k_regimes):
                offset = 0
                indices = []
                for k, v in self.relative_index_regime_purpose[j].items():
                    v = (np.r_[v] + offset).tolist()
                    self.index_regime_purpose[j][k] = v
                    indices.append(v)
                    offset += self.k_parameters[k]
                self.index_regime[j] = np.concatenate(indices).astype(int)
        else:
            raise IndexError("Invalid index")


class MarkovSwitching(tsbase.TimeSeriesModel):
    """
    First-order k-regime Markov switching model

    Parameters
    ----------
    endog : array_like
        The endogenous variable.
    k_regimes : int
        The number of regimes.
    order : int, optional
        The order of the model describes the dependence of the likelihood on
        previous regimes. This depends on the model in question and should be
        set appropriately by subclasses.
    exog_tvtp : array_like, optional
        Array of exogenous or lagged variables to use in calculating
        time-varying transition probabilities (TVTP). TVTP is only used if this
        variable is provided. If an intercept is desired, a column of ones must
        be explicitly included in this array.

    Notes
    -----
    This model is new and API stability is not guaranteed, although changes
    will be made in a backwards compatible way if possible.

    References
    ----------
    Kim, Chang-Jin, and Charles R. Nelson. 1999.
    "State-Space Models with Regime Switching:
    Classical and Gibbs-Sampling Approaches with Applications".
    MIT Press Books. The MIT Press.
    """

    def __init__(
        self,
        endog,
        k_regimes,
        order=0,
        exog_tvtp=None,
        exog=None,
        dates=None,
        freq=None,
        missing="none",
    ):

        # Properties
        self.k_regimes = k_regimes
        self.tvtp = exog_tvtp is not None
        # The order of the model may be overridden in subclasses
        self.order = order

        # Exogenous data
        # TODO add checks for exog_tvtp consistent shape and indices
        self.k_tvtp, self.exog_tvtp = prepare_exog(exog_tvtp)

        # Initialize the base model
        super().__init__(endog, exog, dates=dates, freq=freq, missing=missing)

        # Dimensions
        self.nobs = self.endog.shape[0]

        # Sanity checks
        if self.endog.ndim > 1 and self.endog.shape[1] > 1:
            raise ValueError("Must have univariate endogenous data.")
        if self.k_regimes < 2:
            raise ValueError(
                "Markov switching models must have at least two regimes."
            )
        if not (self.exog_tvtp is None or self.exog_tvtp.shape[0] == self.nobs):
            raise ValueError(
                "Time-varying transition probabilities exogenous"
                " array must have the same number of observations"
                " as the endogenous array."
            )

        # Parameters
        self.parameters = MarkovSwitchingParams(self.k_regimes)
        k_transition = self.k_regimes - 1
        if self.tvtp:
            k_transition *= self.k_tvtp
        self.parameters["regime_transition"] = [1] * k_transition

        # Internal model properties: default is steady-state initialization
        self._initialization = "steady-state"
        self._initial_probabilities = None

    @property
    def k_params(self):
        """
        (int) Number of parameters in the model
        """
        return self.parameters.k_params

    def initialize_steady_state(self):
        """
        Set initialization of regime probabilities to be steady-state values

        Notes
        -----
        Only valid if there are not time-varying transition probabilities.
        """
        if self.tvtp:
            raise ValueError(
                "Cannot use steady-state initialization when"
                " the regime transition matrix is time-varying."
            )

        self._initialization = "steady-state"
        self._initial_probabilities = None

    def initialize_known(self, probabilities, tol=1e-8):
        """
        Set initialization of regime probabilities to use known values
        """
        self._initialization = "known"
        probabilities = np.array(probabilities, ndmin=1)
        if not probabilities.shape == (self.k_regimes,):
            raise ValueError(
                "Initial probabilities must be a vector of shape (k_regimes,)."
            )
        if not np.abs(np.sum(probabilities) - 1) < tol:
            raise ValueError("Initial probabilities vector must sum to one.")
        self._initial_probabilities = probabilities

    def initial_probabilities(self, params, regime_transition=None):
        """
        Retrieve initial probabilities
        """
        params = np.array(params, ndmin=1)
        if self._initialization == "steady-state":
            if regime_transition is None:
                regime_transition = self.regime_transition_matrix(params)
            if regime_transition.ndim == 3:
                regime_transition = regime_transition[..., 0]
            m = regime_transition.shape[0]
            A = np.c_[(np.eye(m) - regime_transition).T, np.ones(m)].T
            try:
                probabilities = np.linalg.pinv(A)[:, -1]
            except np.linalg.LinAlgError as exc:
                raise RuntimeError(
                    "Steady-state probabilities could not be constructed."
                ) from exc
        elif self._initialization == "known":
            probabilities = self._initial_probabilities
        else:
            raise RuntimeError("Invalid initialization method selected.")

        # Slightly bound probabilities away from zero (for filters in log
        # space)
        probabilities = np.maximum(probabilities, 1e-20)

        return probabilities

    def _regime_transition_matrix_tvtp(self, params, exog_tvtp=None):
        if exog_tvtp is None:
            exog_tvtp = self.exog_tvtp
        nobs = len(exog_tvtp)

        regime_transition_matrix = np.zeros(
            (self.k_regimes, self.k_regimes, nobs),
            dtype=np.promote_types(np.float64, params.dtype),
        )

        # Compute the predicted values from the regression
        for i in range(self.k_regimes):
            coeffs = params[self.parameters[i, "regime_transition"]]
            regime_transition_matrix[:-1, i, :] = np.dot(
                exog_tvtp, np.reshape(coeffs, (self.k_regimes - 1, self.k_tvtp)).T
            ).T

        # Perform the logistic transformation
        tmp = np.c_[
            np.zeros((nobs, self.k_regimes, 1)), regime_transition_matrix[:-1, :, :].T
        ].T
        regime_transition_matrix[:-1, :, :] = np.exp(
            regime_transition_matrix[:-1, :, :] - logsumexp(tmp, axis=0)
        )

        # Compute the last column of the transition matrix
        regime_transition_matrix[-1, :, :] = 1 - np.sum(
            regime_transition_matrix[:-1, :, :], axis=0
        )

        return regime_transition_matrix

    def regime_transition_matrix(self, params, exog_tvtp=None):
        """
        Construct the left-stochastic transition matrix

        Notes
        -----
        This matrix will either be shaped (k_regimes, k_regimes, 1) or if there
        are time-varying transition probabilities, it will be shaped
        (k_regimes, k_regimes, nobs).

        The (i,j)th element of this matrix is the probability of transitioning
        from regime j to regime i; thus the previous regime is represented in a
        column and the next regime is represented by a row.

        It is left-stochastic, meaning that each column sums to one (because
        it is certain that from one regime (j) you will transition to *some
        other regime*).
        """
        params = np.array(params, ndmin=1)
        if not self.tvtp:
            regime_transition_matrix = np.zeros(
                (self.k_regimes, self.k_regimes, 1),
                dtype=np.promote_types(np.float64, params.dtype),
            )
            regime_transition_matrix[:-1, :, 0] = np.reshape(
                params[self.parameters["regime_transition"]],
                (self.k_regimes - 1, self.k_regimes),
            )
            regime_transition_matrix[-1, :, 0] = 1 - np.sum(
                regime_transition_matrix[:-1, :, 0], axis=0
            )
        else:
            regime_transition_matrix = self._regime_transition_matrix_tvtp(
                params, exog_tvtp
            )

        return regime_transition_matrix

    def predict(
        self, params, start=None, end=None, probabilities=None, conditional=False
    ):
        """
        In-sample prediction and out-of-sample forecasting

        Parameters
        ----------
        params : ndarray
            Parameters at which to form predictions
        start : int, str, or datetime, optional
            Zero-indexed observation number at which to start forecasting,
            i.e., the first forecast is start. Can also be a date string to
            parse or a datetime type. Default is the the zeroth observation.
        end : int, str, or datetime, optional
            Zero-indexed observation number at which to end forecasting, i.e.,
            the last forecast is end. Can also be a date string to
            parse or a datetime type. However, if the dates index does not
            have a fixed frequency, end must be an integer index if you
            want out of sample prediction. Default is the last observation in
            the sample.
        probabilities : str or array_like, optional
            Specifies the weighting probabilities used in constructing the
            prediction as a weighted average. If a string, can be 'predicted',
            'filtered', or 'smoothed'. Otherwise can be an array of
            probabilities to use. Default is smoothed.
        conditional : bool or int, optional
            Whether or not to return predictions conditional on current or
            past regimes. If False, returns a single vector of weighted
            predictions. If True or 1, returns predictions conditional on the
            current regime. For larger integers, returns predictions
            conditional on the current regime and some number of past regimes.

        Returns
        -------
        predict : ndarray
            Array of out of in-sample predictions and / or out-of-sample
            forecasts.
        """
        if start is None:
            start = self._index[0]

        # Handle start, end
        start, end, out_of_sample, prediction_index = self._get_prediction_index(
            start, end
        )

        if out_of_sample > 0:
            raise NotImplementedError

        # Perform in-sample prediction
        predict = self.predict_conditional(params)
        squeezed = np.squeeze(predict)

        # Check if we need to do weighted averaging
        if squeezed.ndim - 1 > conditional:
            # Determine in-sample weighting probabilities
            if probabilities is None or probabilities == "smoothed":
                results = self.smooth(params, return_raw=True)
                probabilities = results.smoothed_joint_probabilities
            elif probabilities == "filtered":
                results = self.filter(params, return_raw=True)
                probabilities = results.filtered_joint_probabilities
            elif probabilities == "predicted":
                results = self.filter(params, return_raw=True)
                probabilities = results.predicted_joint_probabilities

            # Compute weighted average
            predict = predict * probabilities
            for _ in range(predict.ndim - 1 - int(conditional)):
                predict = np.sum(predict, axis=-2)
        else:
            predict = squeezed

        return predict[start : end + out_of_sample + 1]

    def predict_conditional(self, params):
        """
        In-sample prediction, conditional on the current, and possibly past,
        regimes

        Parameters
        ----------
        params : array_like
            Array of parameters at which to perform prediction.

        Returns
        -------
        predict : array_like
            Array of predictions conditional on current, and possibly past,
            regimes
        """
        raise NotImplementedError

    def _conditional_loglikelihoods(self, params):
        """
        Compute likelihoods conditional on the current period's regime (and
        the last self.order periods' regimes if self.order > 0).

        Must be implemented in subclasses.
        """
        raise NotImplementedError

    def _filter(self, params, regime_transition=None):
        # Get the regime transition matrix if not provided
        if regime_transition is None:
            regime_transition = self.regime_transition_matrix(params)
        # Get the initial probabilities
        initial_probabilities = self.initial_probabilities(params, regime_transition)

        # Compute the conditional likelihoods
        conditional_loglikelihoods = self._conditional_loglikelihoods(params)

        # Apply the filter
        return (
            regime_transition,
            initial_probabilities,
            conditional_loglikelihoods,
        ) + cy_hamilton_filter_log(
            initial_probabilities,
            regime_transition,
            conditional_loglikelihoods,
            self.order,
        )

    def filter(
        self,
        params,
        transformed=True,
        cov_type=None,
        cov_kwds=None,
        return_raw=False,
        results_class=None,
        results_wrapper_class=None,
    ):
        """
        Apply the Hamilton filter

        Parameters
        ----------
        params : array_like
            Array of parameters at which to perform filtering.
        transformed : bool, optional
            Whether or not `params` is already transformed. Default is True.
        cov_type : str, optional
            See `fit` for a description of covariance matrix types
            for results object.
        cov_kwds : dict or None, optional
            See `fit` for a description of required keywords for alternative
            covariance estimators
        return_raw : bool,optional
            Whether or not to return only the raw Hamilton filter output or a
            full results object. Default is to return a full results object.
        results_class : type, optional
            A results class to instantiate rather than
            `MarkovSwitchingResults`. Usually only used internally by
            subclasses.
        results_wrapper_class : type, optional
            A results wrapper class to instantiate rather than
            `MarkovSwitchingResults`. Usually only used internally by
            subclasses.

        Returns
        -------
        MarkovSwitchingResults
        """
        params = np.array(params, ndmin=1)

        if not transformed:
            params = self.transform_params(params)

        # Save the parameter names
        self.data.param_names = self.param_names

        # Get the result
        names = [
            "regime_transition",
            "initial_probabilities",
            "conditional_loglikelihoods",
            "filtered_marginal_probabilities",
            "predicted_joint_probabilities",
            "joint_loglikelihoods",
            "filtered_joint_probabilities",
            "predicted_joint_probabilities_log",
