# -*- coding: utf-8 -*-
from abc import abstractmethod
from typing import Sequence

from scipy.linalg import block_diag

from ..base import Model, GaussianModel
from ...base import Property


class TransitionModel(Model):
    """Transition Model base class"""

    @property
    def ndim(self):
        return self.ndim_state

    @property
    @abstractmethod
    def ndim_state(self):
        """Number of state dimensions"""
        pass


class _CombinedGaussianTransitionModel(TransitionModel, GaussianModel):
    model_list: Sequence[GaussianModel] = Property(doc="List of Transition Models.")

    @property
    def ndim_state(self):
        """ndim_state getter method

        Returns
        -------
        : :class:`int`
            The number of combined model state dimensions.
        """
        return sum(model.ndim_state for model in self.model_list)

    def covar(self, **kwargs):
        """Returns the transition model noise covariance matrix.

        Returns
        -------
        : :class:`stonesoup.types.state.CovarianceMatrix` of shape\
        (:py:attr:`~ndim_state`, :py:attr:`~ndim_state`)
            The process noise covariance.
        """

        covar_list = [model.covar(**kwargs) for model in self.model_list]
        return block_diag(*covar_list)
