"""
Author: Mingjian He <mh1@stanford.edu>

ssm module contains a general state-space model class used in SOMATA
"""

from somata.exact_inference import kalman, djkalman, inverse
from somata.utils import estimate_r
import numpy as np
import numbers
from collections.abc import Iterable
from copy import deepcopy
from joblib import Parallel, delayed, cpu_count
from scipy.linalg import block_diag


class StateSpaceModel(object):
    """ StateSpaceModel is a general object class for state-space modeling in SOMATA """
    # Class attributes (should not be manually changed or be mutable)
    # only listed here as a reference list for all attributes under Ssm

    # Attributes are intentionally not implemented as properties with
    # decorators in order to allow for flexible input types. Type-hints
    # are omitted for the same reason. numpy.asanyarray() will handle
    # type conversions in the end.

    type = 'ssm'
    nstate = None
    ncomp = None
    nchannel = None
    ntime = None
    nmodel = None
    components = None
    comp_nstates = None
    F = None
    Q = None
    mu0 = None
    S0 = None
    G = None
    R = None
    y = None
    Fs = None
    stackable = ('F', 'Q', 'mu0', 'S0', 'G', 'R')
    default_G = None

    def __init__(self, components=None, F=None, Q=None, mu0=None, S0=None, G=None, R=None, y=None, Fs=None):
        """
        Constructor method for StateSpaceModel class
        :param components: a list of independent components in the model
        :param F: transition matrix
        :param Q: state noise covariance matrix
        :param mu0: initial state mean vector
        :param S0: initial state covariance matrix
        :param G: observation matrix (row major)
        :param R: observation noise covariance matrix
        :param y: observed data (row major, can be multivariate)
        :param Fs: sampling frequency in Hz
        """
        # Instance attributes
        self.F = self._process_constructor_input(F)
        self.Q = self._process_constructor_input(Q)
        self.mu0 = self._process_constructor_input(mu0)
        self.S0 = self._process_constructor_input(S0)
        self.y = self._must_be_row(self._process_constructor_input(y))
        self.Fs = np.float64(Fs) if Fs is not None else None

        # Initialize components
        if components is None:
            self.ncomp = 0
            self.components = None
            self.comp_nstates = [0]
        elif isinstance(components, str):
            self._auto_populate_components(components)
        else:
            self._initialize_from_components(components)

        # Initialize observation matrix and noise
        self._initialize_observation_matrix(G)
        self._initialize_observation_noise(R)

        # Check all dimensions
        self._check_dimensions()
        self._check_model_stack()

    # Dunder methods - magic methods and arithmetic-like operations
    def __repr__(self):
        """ Unambiguous and concise representation when calling StateSpaceModel() """
        return 'Ssm(' + str(self.ncomp) + ')<' + hex(id(self))[-4:] + '>'

    def __str__(self):
        """ Helpful information when calling print(StateSpaceModel()) """
        print_str = "<Ssm object at " + hex(id(self)) + ">\n " + \
                    "{0:8} = {1: <5} ".format("nstate", str(self.nstate)) + \
                    "{0:8} = {1}\n ".format("ncomp", str(self.ncomp)) + \
                    "{0:8} = {1: <5} ".format("nchannel", str(self.nchannel)) + \
                    "{0:8} = {1}\n ".format("ntime", str(self.ntime)) + \
                    "{0:8} = {1}\n ".format("nmodel", str(self.nmodel)) + \
                    "components = {}\n ".format(str(self.components)) + \
                    "{0:3}.shape = {1: <10} ".format("F", [str(x.shape) if x is not None else 'None' for x
                                                           in [self.F]][0]) + \
                    "{0:3}.shape = {1}\n ".format("Q", [str(x.shape) if x is not None else 'None' for x
                                                        in [self.Q]][0]) + \
                    "{0:3}.shape = {1: <10} ".format("mu0", [str(x.shape) if x is not None else 'None' for x
                                                             in [self.mu0]][0]) + \
                    "{0:3}.shape = {1}\n ".format("S0", [str(x.shape) if x is not None else 'None' for x
                                                         in [self.S0]][0]) + \
                    "{0:3}.shape = {1: <10} ".format("G", [str(x.shape) if x is not None else 'None' for x
                                                           in [self.G]][0]) + \
                    "{0:3}.shape = {1}\n ".format("R", [str(x.shape) if x is not None else 'None' for x
                                                        in [self.R]][0]) + \
                    "{0:3}.shape = {1: <10} ".format("y", [str(x.shape) if x is not None else 'None' for x
                                                           in [self.y]][0]) + \
                    "Fs = {0}\n ".format([str(x) + ' Hz' if x is not None else 'None' for x in [self.Fs]][0])
        return print_str

    def ssm(self):
        """ Create a StateSpaceModel object with the same attributes as the OscillatorModel object """
        return StateSpaceModel(components=self.components, F=self.F, Q=self.Q, mu0=self.mu0, S0=self.S0,
                               G=self.G, R=self.R, y=self.y, Fs=self.Fs)

    def __len__(self):
        return self.nmodel

    def __add__(self, other):
        return self._stack(other)

    def __radd__(self, other):
        return self.__add__(other)

    def __mul__(self, other):
        return self._permute(other)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __matmul__(self, other):
        return self.__mul__(other)

    def __rmatmul__(self, other):
        return self.__matmul__(other)

    def __pow__(self, other: int):
        if other == 0:
            return StateSpaceModel()
        elif other > 0:
            # simply a magic method to call _expand_attr() with deepcopy()
            out_copy = deepcopy(self)
            out_copy._expand_attr()
            return out_copy
        else:
            raise SyntaxError('StateSpaceModel objects cannot be taken to negative power.')

    def __rpow__(self, other: int):
        return self.__pow__(other)

    # Check methods - sanity checks to ensure model validity
    def _check_dimensions(self, check_scope=(0, 1, 2)):
        """
        Check dimensions of instance attributes, ignoring None attributes
        :param check_scope: (int) specifying which ones to check.
            0 = state dimensions,
            1 = observed data dimensions,
            2 = components and observation matrix
        """
        if 0 in check_scope:
            # check the dimension of states
            check_state_tuple = (self.F, self.Q, self.mu0, self.S0)
            check_state_n = len(check_state_tuple)
            check_state_result = np.zeros(check_state_n, dtype=np.int_)
            for ii in range(check_state_n):
                if check_state_tuple[ii] is not None:
                    attr_shape = check_state_tuple[ii].shape
                    check_state_result[ii] = attr_shape[0]
                    if ii != 2:  # mu0 is a vector therefore no need to check
                        assert attr_shape[0] == attr_shape[1], 'Matrix attribute is not square!'

            nstates = np.unique(check_state_result[np.nonzero(check_state_result)])
            if len(nstates) == 0:
                self.nstate = 0
            else:
                assert len(nstates) == 1, 'Number of states is ambiguous and inconsistent across attributes.'
                self.nstate = int(nstates[0])

        if 1 in check_scope:
            # check the dimension of observation channels
            self.nchannel = 0
            self.ntime = 0
            if self.R is not None:
                assert self.R.shape[0] == self.R.shape[1], 'Matrix attribute is not square!'
                self.nchannel = self.R.shape[0]
            if self.y is not None:
                self.nchannel = self.y.shape[0]
                self.ntime = self.y.shape[1]
                if self.R is not None:
                    assert self.y.shape[0] == self.R.shape[0], \
                        'Observation dimensions are inconsistent across attributes.'

        if 2 in check_scope:
            # check the dimensions of components and observation matrix
            if self.components is not None:
                assert self.nstate == sum([x.default_G.shape[1] for x in self.components]), \
                    'Input components give a different state dimension from that of other attributes.'
                assert self.ncomp == len(self.components), 'Numbers of components are inconsistent.'
                self.comp_nstates = [x.default_G.shape[1] for x in self.components]
            else:
                self.comp_nstates = [0]

            if self.G is not None:
                if self.nchannel > 0:
                    assert self.G.shape[0] == self.nchannel, 'Observation matrix dimension mismatches nchannel.'
                if self.nstate > 0:
                    assert self.G.shape[1] == self.nstate, 'Observation matrix dimension mismatches nstate.'

    def _check_model_stack(self):
        """
        Check the number of alternative models stacked in instance attributes.
        third axis is reserved for stacking different models
        """
        # Different components and y should reside in different instances, therefore not stackable
        assert self.components is None or isinstance(self.components, list), \
            'Component structure should not be stacked in a single StateSpaceModel instance.'
        assert self.y is None or len(self.y.shape) < 3, \
            'Observed data y should not be stacked in a single StateSpaceModel instance.'
        assert self.Fs is None or isinstance(self.Fs, numbers.Number), \
            'Sampling frequency is not a valid Number type.'

        # The rest of instance attributes can be stacked with different models
        stack_numbers = self._get_stack_numbers()
        model_stack = np.unique(stack_numbers[stack_numbers != 1])
        if len(model_stack) == 0:
            self.nmodel = 1
        else:
            assert len(model_stack) == 1, 'More than one model stack numbers. Invoke __mul__ if trying to permute.'
            self.nmodel = int(model_stack[0])

    def _check_observed_data(self, other):
        """ Used to check if two objects have the same observed data """
        if self.y is not None and other.y is not None:
            assert self.y.shape == other.y.shape, 'Observed data have different dimensions.'
            # randomly select one time point to spot check
            rindex = np.random.randint(self.y.shape[0])
            cindex = np.random.randint(self.y.shape[1])
            assert self.y[rindex, cindex] == other.y[rindex, cindex], 'Observed data are not identical.'
        if self.Fs is not None and other.Fs is not None:
            assert self.Fs == other.Fs, 'Observed data have different sampling frequencies.'

    # Syntactic sugar methods - useful methods to make manipulations easier
    def _get_stack_numbers(self):
        """ Obtain the number of models stacked in each attribute """
        check_stack_tuple = tuple([getattr(self, x) for x in self.stackable])
        check_stack_n = len(check_stack_tuple)
        stack_numbers = np.ones(check_stack_n, dtype=np.int_)
        for ii in range(check_stack_n):
            if check_stack_tuple[ii] is not None:
                attr_shape = check_stack_tuple[ii].shape
                if len(attr_shape) > 2:
                    stack_numbers[ii] = attr_shape[2]
        return stack_numbers

    def _stack(self, other):
        """ Stack attributes in two objects """
        self._check_observed_data(other)
        new_obj = self._concat_attr(other)
        new_obj._check_model_stack()
        return new_obj

    def _permute(self, other):
        """ Permute attributes in two objects """
        self._check_observed_data(other)
        new_obj = self._concat_attr(other)
        new_obj._expand_attr()
        return new_obj

    def _concat_attr(self, other, attrs=None, ignore_duplicate=True):
        """
        Concatenate two objects via stacking new model attributes,
        duplicate checking is enabled; therefore repeated attributes
        will be ignored by turning add_flag to False, unless the
        ignore_duplicate flag is set to False
        """
        new_obj = deepcopy(self)
        if attrs is None:
            attrs = ('components', 'y', 'Fs') + self.stackable
        elif not isinstance(attrs, tuple):
            attrs = tuple(attrs)

        for attr_name in attrs:
            self_attr = getattr(self, attr_name)
            other_attr = getattr(other, attr_name)
            if self_attr is None and other_attr is None:
                pass
            elif other_attr is None:
                pass
            elif self_attr is None:
                setattr(new_obj, attr_name, deepcopy(other_attr))
            else:
                if attr_name == 'Fs':
                    assert self_attr == other_attr, 'Observed data have different sampling frequencies.'
                if attr_name == 'components':
                    assert len(self_attr) == len(other_attr), 'Components have different shapes.'
                elif len(other_attr.shape) > 2:
                    for j in range(other_attr.shape[2]):
                        model_attr = other_attr[:, :, j]
                        add_flag = True
                        if len(self_attr.shape) > 2:
                            for k in range(self_attr.shape[2]):
                                if (self_attr[:, :, k] == model_attr).all():  # type: ignore
                                    add_flag = False
                        else:
                            if (self_attr == model_attr).all():  # type: ignore
                                add_flag = False
                        if add_flag or ignore_duplicate is False:
                            setattr(new_obj, attr_name, np.dstack([getattr(new_obj, attr_name), model_attr]))
                else:
                    model_attr = other_attr
                    add_flag = True
                    if len(self_attr.shape) > 2:
                        for k in range(self_attr.shape[2]):
                            if (self_attr[:, :, k] == model_attr).all():  # type: ignore
                                add_flag = False
                    else:
                        if (self_attr == model_attr).all():  # type: ignore
                            add_flag = False
                    if add_flag or ignore_duplicate is False:
                        setattr(new_obj, attr_name, np.dstack([getattr(new_obj, attr_name), model_attr]))
        new_obj._check_dimensions()
        return new_obj

    def _expand_attr(self):
        """ Expand stacked attributes in an object """
        stack_numbers = self._get_stack_numbers()
        update_idx = [x for x in range(len(stack_numbers)) if (stack_numbers > 1)[x]]
        new_stack_number = int(stack_numbers[update_idx].prod())
        # expand with permutation of attributes to form stacked models
        last_attr_multiplier = 1
        for j in range(len(update_idx)):
            attr_name = self.stackable[update_idx[j]]
            current_attr = getattr(self, attr_name)
            rep_num = int(new_stack_number // last_attr_multiplier // stack_numbers[update_idx[j]])
            temp_attr = np.dstack([np.dstack([current_attr[:, :, n] for _ in range(rep_num)])
                                   for n in range(current_attr.shape[2])])
            setattr(self, attr_name, np.dstack([temp_attr for _ in range(last_attr_multiplier)]))
            last_attr_multiplier = int(last_attr_multiplier * stack_numbers[update_idx[j]])
        self._check_dimensions()
        self._check_model_stack()

    def _auto_populate_components(self, type_str):
        """ Auto-populate components with a given type of components """
        assert isinstance(type_str, str), 'components should be a string in this method.'
        self._check_dimensions((0,))  # fill in the nstate attribute

        if type_str == 'Osc':  # auto-populate with oscillator components
            self._auto_populate_osc_components()
        elif type_str == 'Arn':
            self._auto_populate_arn_components()
        elif type_str == 'Gen':
            self._auto_populate_gen_components()
        else:
            raise ValueError('Specified component type is not valid.')

        self._check_dimensions((2,))  # fill in the comp_nstates attribute

    def _auto_populate_osc_components(self):
        """ Initialize all components to be Matsuda oscillators during constructor """
        try:  # this has to stay here to avoid circular import before class definitions
            from .osc import OscillatorModel
        except (ImportError, ModuleNotFoundError):
            from somata import OscillatorModel

        assert (self.nstate % 2) == 0, 'Cannot default to Osc components with odd nstate.'
        self.ncomp = self.nstate // 2
        if self.ncomp > 0:
            self.components = [OscillatorModel() for _ in range(self.ncomp)]
        else:
            self.components = None

    def _auto_populate_arn_components(self):
        """ Initialize all components to be autoregressive models during constructor """
        try:  # this has to stay here to avoid circular import before class definitions
            from .arn import AutoRegModel
        except (ImportError, ModuleNotFoundError):
            from somata import AutoRegModel

        if self.F is not None:
            # Guess how many Arn components are contained in the model parameters
            if self.Q is not None:
                assert np.all(self.Q[~np.eye(self.Q.shape[0], dtype=bool)] == 0), \
                    'Found non-zero non-diagonal elements when trying to auto-populate Arn components.'
                Q_diag = np.diag(self.Q)
                comp_nstates = np.diff(np.append(np.nonzero(Q_diag), len(Q_diag)))
            else:
                comp_nstates = AutoRegModel._guess_ar_order(self, F=self.F)
                Q_diag = np.zeros(sum(comp_nstates))
                Q_diag[np.cumsum(np.hstack(([0], comp_nstates[:-1])))] = AutoRegModel.default_sigma2
                self.Q = np.diag(Q_diag)  # fill in self.Q with default sigma2

            # Make sure the state dimensions are correct
            assert sum(comp_nstates) == self.nstate, 'Arn components do not have the correct total number of states.'
            self.ncomp = len(comp_nstates)

            # Create the autoregressive model components
            self.components = []
            for n in range(len(comp_nstates)):
                start_idx = sum(comp_nstates[:n])
                end_idx = sum(comp_nstates[:n+1])
                coeff = self.F[start_idx, start_idx:end_idx]
                sigma2 = self.Q[start_idx, start_idx]
                mu0 = self.mu0[start_idx:end_idx] if self.mu0 is not None else None
                S0 = self.Q[start_idx, start_idx] if self.S0 is not None else None
                self.components.append(AutoRegModel(coeff=coeff, sigma2=sigma2, mu0=mu0, S0=S0))
        else:
            assert self.Q is None, 'No F provided but state noise covariance input is given.'
            self.ncomp = 0
            self.components = None

    def _auto_populate_gen_components(self):
        """ Initialize with a single general state-space model during constructor """
        try:  # this has to stay here to avoid circular import before class definitions
            from .gen import GeneralSSModel
        except (ImportError, ModuleNotFoundError):
            from somata import GeneralSSModel

        if self.nstate > 0:
            self.ncomp = 1
            self.components = [GeneralSSModel()]
            self.components[0].default_G = np.ones((1, self.nstate), dtype=np.float64)
        else:
            self.ncomp = 0
            self.components = None

    def _initialize_from_components(self, components):
        """ Initialize the StateSpaceModel instance from given components during constructor """
        assert not isinstance(components, str), 'components should not be a string in this method.'

        # Form a list of single components without nesting
        components = list(components) if isinstance(components, Iterable) else [components]
        comp_list = []
        for n in range(len(components)):
            current_component = components[n]
            if current_component.ncomp <= 1:
                comp_list.append(current_component)
            else:
                assert isinstance(current_component.components, list), 'Invalid components attribute.'
                assert (np.asarray([x.ncomp for x in current_component.components]) <= 1).all(), \
                    'Components should not have nested multiple components in them.'
                comp_list += current_component.components

        self.ncomp = len(comp_list)
        self.components = comp_list
        self.comp_nstates = [x.default_G.shape[1] for x in self.components]

        # Use the Ssm.concat_() method to concatenate parameters from components in order to fill them in
        concat_model = components[0]
        for n in range(1, len(components)):
            concat_model = StateSpaceModel.concat_(concat_model, components[n], skip_components=True)

        # Fill in model parameters if available from components
        self._setattr_when_not_none(concat_model, ('F', 'Q', 'mu0', 'S0', 'G', 'R', 'Fs'))

        # Fill in observed data if available from components
        self._check_observed_data(concat_model)
        self.y = self._return_not_none(self.y, concat_model.y)

    def _setattr_when_not_none(self, other, attrs):
        """ Update an instance attribute when the other model has non-empty values """
        attrs = tuple(attrs) if not isinstance(attrs, tuple) else attrs
        for attr_name in attrs:
            self_attr = getattr(self, attr_name)
            other_attr = getattr(other, attr_name)
            if other_attr is not None:
                if self_attr is None:
                    setattr(self, attr_name, other_attr)
                else:
                    assert np.all(self_attr == other_attr), \
                        'Incompatible attribute ' + attr_name + ' when initializing from components.'

    def _initialize_observation_matrix(self, G):
        """ Initialize the observation matrix instance attribute during constructor """
        G = self.G if G is None else G
        if G is None:
            if self.ncomp > 0:
                nchannel = self.y.shape[0] if self.y is not None else 1  # default to one observation channel
                self.G = np.tile(self._must_be_row(self._process_constructor_input(
                    np.hstack([x.default_G for x in self.components]))), (nchannel, 1))
            else:
                self.G = None
        else:
            self.G = self._must_be_row(self._process_constructor_input(G))

    def _initialize_observation_noise(self, R):
        """ Initialize the observation noise covariance matrix during constructor """
        # Update R if not provided but can be estimated from y and Fs
        R = self.R if R is None else R
        if R is None and self.y is not None and self.Fs is not None:
            R = estimate_r(y=self.y, Fs=self.Fs, freq_cutoff=self.Fs / 2 - 20)

        self.R = self._process_constructor_input(R)

    def stack_attr(self, attr_name, new_attr):
        """ Stack new_attr into the third axis of an instance attribute """
        if attr_name == 'components' or attr_name == 'y' or attr_name == 'Fs':
            raise AttributeError(attr_name + " shouldn't be stacked.")
        # using _concat_attr can automatically ignore duplicated attributes
        temp_obj = self._concat_attr(StateSpaceModel(**{attr_name: new_attr}), attrs=attr_name)
        setattr(self, attr_name, getattr(temp_obj, attr_name))
        self._check_dimensions()
        try:
            self._check_model_stack()
        except AssertionError:
            print("WARNING: nmodel cannot be updated automatically due to ambiguous model stack numbers.")

    def stack_to_array(self):
        """
        Unstack attributes into an array of objects,
        which is duck-typing-equivalent with a list
        """
        self._check_model_stack()
        stack_numbers = self._get_stack_numbers()
        update_idx = [x for x in range(len(stack_numbers)) if (stack_numbers > 1)[x]]
        ssm_array = np.empty(self.nmodel, dtype=StateSpaceModel)  # mutable array
        for ii in range(self.nmodel):
            temp_obj = deepcopy(self)
            for j in range(len(update_idx)):
                temp_attr = deepcopy(getattr(self, self.stackable[update_idx[j]]))
                setattr(temp_obj, self.stackable[update_idx[j]], temp_attr[:, :, ii])
            temp_obj._check_dimensions()
            temp_obj._check_model_stack()
            ssm_array[ii] = temp_obj
        return ssm_array

    def copy(self, drop_y=False):
        """ Make a copy of the StateSpaceModel Object """
        self_copy = deepcopy(self)
        if drop_y:
            self_copy.y = None
        return self_copy

    def append(self, other):
        """
        Append an object to an existing StateSpaceModel object,
        by calling concat_() and copying over attributes
        """
        self._setattr_for_append(temp_obj=self.concat_(other))

    def rappend(self, other):
        """
        Append an existing StateSpaceModel object to an object,
        positioned on the right side, i.e., ohter is appended from
        the left, by calling concat_() and copying over attributes
        """
        self._setattr_for_append(temp_obj=other.concat_(self))

    def _setattr_for_append(self, temp_obj):
        """
        Copy over attributes from temp_obj to self for appending.
        If self and other are of the same type, modify attributes
        in place, otherwise, raise an exception
        """
        if self.type == temp_obj.type:
            for attr in self.__dict__.keys():
                setattr(self, attr, getattr(temp_obj, attr))
        else:
            raise TypeError('Cannot append two SOMATA models of different classes. Use concat_() instead.')

    def concat_(self, other, skip_components=False):
        """
        Join two StateSpaceModel objects together by concatenating the
        components, and return a new object with the new total number
        of components. The mutable attribute components has new memory
        addresses for each component
        """
        # If both have observed data, they should be identical
        self._check_observed_data(other)
        y = self._return_not_none(self.y, other.y)
        Fs = self._return_not_none(self.Fs, other.Fs)
        if self.R is not None and other.R is not None:
            assert (self.R == other.R).all(), 'Cannot concatenate two objects with conflicting R.'  # type: ignore
            R = self.R
        else:
            R = self._return_not_none(self.R, other.R)

        # Fill components with parameters if concatenating with different classes of SOMATA basic models
        if self.type != other.type and not skip_components:
            components_prefill_self = self.fill_components()
            components_prefill_other = other.fill_components()

        # Set up the new components attribute using deepcopy(), without maintaining memory addresses
        new_comp = [None] * (self.ncomp + other.ncomp)
        if len(new_comp) == 0 or skip_components:
            new_comp = None
        else:
            new_comp[:self.ncomp] = [None] * self.ncomp if self.components is None else deepcopy(self.components)
            new_comp[self.ncomp:] = [None] * other.ncomp if other.components is None else deepcopy(other.components)

        # Revert any changes made to the components attribute of the original objects
        if self.type != other.type and not skip_components:
            self.unfill_components(components_prefill_self)
            other.unfill_components(components_prefill_other)

        # Configure the rest of attributes that are immutable
        # F
        if self.F is None and other.F is None:
            F = None
        elif self.F is None:
            F = block_diag(np.empty((self.nstate, self.nstate), dtype=object), other.F)
        elif other.F is None:
            F = block_diag(self.F, np.empty((other.nstate, other.nstate), dtype=object))
        else:
            F = block_diag(self.F, other.F)

        # Q
        if self.Q is None and other.Q is None:
            Q = None
        elif self.Q is None:
            Q = block_diag(np.empty((self.nstate, self.nstate), dtype=object), other.Q)
        elif other.Q is None:
            Q = block_diag(self.Q, np.empty((other.nstate, other.nstate), dtype=object))
        else:
            Q = block_diag(self.Q, other.Q)

        # mu0
        if self.mu0 is None and other.mu0 is None:
            mu0 = None
        elif self.mu0 is None:
            mu0 = np.vstack([np.zeros((self.nstate, 1), dtype=np.float64), other.mu0])
        elif other.mu0 is None:
            mu0 = np.vstack([self.mu0, np.zeros((other.nstate, 1), dtype=np.float64)])
        else:
            mu0 = np.vstack([self.mu0, other.mu0])

        # S0
        if self.S0 is None and other.S0 is None:
            S0 = None
        elif self.S0 is None:
            S0 = block_diag(self.Q, other.S0) if self.Q is not None else \
                block_diag(np.empty((self.nstate, self.nstate), dtype=object), other.S0)
        elif other.mu0 is None:
            S0 = block_diag(self.S0, other.Q) if other.Q is not None else \
                block_diag(self.S0, np.empty((other.nstate, other.nstate), dtype=object))
        else:
            S0 = block_diag(self.S0, other.S0)

        # G
        if self.G is None and other.G is None:
            G = None
        else:
            nchannel = y.shape[0] if y is not None else 1  # default to one observation channel

            if self.G is None:
                assert self.ncomp > 0, 'Cannot concatenate when self has no components.'
                G_self = np.tile(self._must_be_row(self._process_constructor_input(
                        np.hstack([x.default_G for x in self.components]))), (nchannel, 1))
            else:
                G_self = self.G

            if other.G is None:
                assert other.ncomp > 0, 'Cannot concatenate when other has no components.'
                G_other = np.tile(other._must_be_row(other._process_constructor_input(
                    np.hstack([x.default_G for x in other.components]))), (nchannel, 1))
            else:
                G_other = other.G

            G = np.hstack([G_self, G_other])

        return StateSpaceModel(components=new_comp, F=F, Q=Q, mu0=mu0, S0=S0, G=G, R=R, y=y, Fs=Fs)

    def remove_component(self, comp_idx):
        """ Remove a component from the Ssm object """
        if isinstance(comp_idx, list):  # recursive case
            comp_idx_list = list(np.sort(comp_idx)[::-1])
            for comp_idx in comp_idx_list:
                self.remove_component(comp_idx=comp_idx)
        else:  # base case
            start_idx = sum(self.comp_nstates[:comp_idx])
            end_idx = sum(self.comp_nstates[:comp_idx+1])
            slice_idx = range(start_idx, end_idx)
            self.nstate -= self.comp_nstates[comp_idx]
            self.ncomp -= 1
            _ = self.comp_nstates.pop(comp_idx)
            _ = self.components.pop(comp_idx)
            if self.F is not None:
                self.F = np.delete(self.F, slice_idx, axis=0)
                self.F = np.delete(self.F, slice_idx, axis=1)
            if self.Q is not None:
                self.Q = np.delete(self.Q, slice_idx, axis=0)
                self.Q = np.delete(self.Q, slice_idx, axis=1)
            if self.mu0 is not None:
                self.mu0 = np.delete(self.mu0, slice_idx, axis=0)
            if self.S0 is not None:
                self.S0 = np.delete(self.S0, slice_idx, axis=0)
                self.S0 = np.delete(self.S0, slice_idx, axis=1)
            if self.G is not None:
                self.G = np.delete(self.G, slice_idx, axis=1)

            # Set attributes to default values if nothing left
            self.comp_nstates = 0 if len(self.comp_nstates) == 0 else self.comp_nstates
            self.components = None if len(self.components) == 0 else self.components
            self.F = None if len(self.F) == 0 else self.F
            self.Q = None if len(self.Q) == 0 else self.Q
            self.mu0 = None if len(self.mu0) == 0 else self.mu0
            self.S0 = None if len(self.S0) == 0 else self.S0

    def fill_components(self, empty_comp=None, deep_copy=True):
        """
        Fill components with attribute content with deepcopy(),
        observed data y is copied using memory address to save space,
        so it remains mutable (be careful when you slice-alter y)
        """
        # Create an empty component instance
        empty_comp = empty_comp or StateSpaceModel()

        # Save the current components in case we need to revert to unfilled status
        components_prefill = deepcopy(self.components)

        for ii in range(self.ncomp):
            current_component: StateSpaceModel = self.components[ii]  # typehint to superclass to be flexible
            start_idx = sum(self.comp_nstates[:ii])
            end_idx = sum(self.comp_nstates[:ii+1])

            # Provide an empty component in the components attribute
            setattr(current_component, 'components', [deepcopy(empty_comp)])
            current_component.ncomp = 1
            current_component.components[0].default_G = current_component.default_G

            if getattr(self, 'F') is not None:
                if deep_copy:
                    F = deepcopy(getattr(self, 'F')[start_idx:end_idx, start_idx:end_idx])
                else:
                    F = getattr(self, 'F')[start_idx:end_idx, start_idx:end_idx]
                setattr(current_component, 'F', F)
            if getattr(self, 'Q') is not None:
                if deep_copy:
                    Q = deepcopy(getattr(self, 'Q')[start_idx:end_idx, start_idx:end_idx])
                else:
                    Q = getattr(self, 'Q')[start_idx:end_idx, start_idx:end_idx]
                setattr(current_component, 'Q', Q)
            if getattr(self, 'mu0') is not None:
                if deep_copy:
                    mu0 = deepcopy(getattr(self, 'mu0')[start_idx:end_idx])
                else:
                    mu0 = getattr(self, 'mu0')[start_idx:end_idx]
                setattr(current_component, 'mu0', mu0)
            if getattr(self, 'S0') is not None:
                if deep_copy:
                    S0 = deepcopy(getattr(self, 'S0')[start_idx:end_idx, start_idx:end_idx])
                else:
                    S0 = getattr(self, 'S0')[start_idx:end_idx, start_idx:end_idx]
                setattr(current_component, 'S0', S0)
            if getattr(self, 'G') is not None:
                if deep_copy:
                    G = deepcopy(getattr(self, 'G')[:, start_idx:end_idx])
                else:
                    G = getattr(self, 'G')[:, start_idx:end_idx]
                setattr(current_component, 'G', G)

            # R remains unchanged across components
            if deep_copy:
                setattr(current_component, 'R', deepcopy(getattr(self, 'R')))
            else:
                setattr(current_component, 'R', getattr(self, 'R'))

            setattr(current_component, 'y', getattr(self, 'y'))  # point to the same memory address
            setattr(current_component, 'Fs', getattr(self, 'Fs'))  # Fs should have immutable datatype

            # Use check methods in component subclasses to update dimension attributes
            current_component._check_dimensions()
            current_component._check_model_stack()

            # Update component parameters
            current_component.update_comp_param()

        return components_prefill

    def unfill_components(self, components_prefill):
        """ Unfill components to a prefill state """
        assert len(self.components) == len(components_prefill), 'Different numbers of components during unfill.'
        for ii in range(self.ncomp):
            attr_dict = components_prefill[ii].__dict__
            for attr in attr_dict.keys():
                setattr(self.components[ii], attr, getattr(components_prefill[ii], attr))

    def get_default_q(self, components=None, E=None):
        """
        Get the default structure of state noise covariance
        matrix Q in the Q_basis block diagonal form
        """
        components = self.components if components is None else components
        default_Q = block_diag(*[x.get_default_q(components=x, E=E) for x in components])
        return default_Q

    def simulate(self, duration=10):
        """
        Simulate observed data from a state-space model.

        Parameters
        ----------
        self : StateSpaceModel
            A SOMATA state-space model instance.
        duration : int
            Simulation duration in seconds. Default: 10 seconds.

        Returns
        -------
        x : ndarray
            Hidden states.
        y : ndarray
            Simulated observed data.
        """
        # total number of samples in the simulated time series
        T = int(duration * self.Fs)

        # initialize tallies for the hidden states and observations
        x = np.zeros((self.nstate, T + 1))  # one sample longer than y
        y = np.zeros((self.nchannel, T))

        # initial state at t=0
        x[:, 0] = np.random.multivariate_normal(np.squeeze(self.mu0), self.S0)

        # iterate through the rest of time points and generate observations
        for ii in range(1, T + 1):
            x[:, ii] = self.F @ x[:, ii - 1] + np.random.multivariate_normal(np.zeros(self.nstate), self.Q)
            y[:, ii - 1] = self.G @ x[:, ii] + np.random.normal(np.zeros(self.nchannel), self.R)

        return x, y

    @staticmethod
    def _process_constructor_input(a):
        """
        Process inputs to the constructor of Ssm class.
        third axis is reserved for stacking different models instantiated with tuple inputs
        """
        if a is None:
            return a
        elif isinstance(a, tuple):
            return np.asanyarray(np.dstack(a), dtype=np.float64)
        else:
            # always break the link to input data memory address
            a = deepcopy(np.asanyarray(a, dtype=np.float64))

            if len(a.shape) == 0:  # always promote to (r,c) 2D arrays
                return a[None, None]
            elif len(a.shape) == 1:
                return a[:, None]
            else:
                return a

    @staticmethod
    def _must_be_row(x):
        """ Ensure input vector is a row vector """
        if x is None:
            return x
        elif x.shape[1] == 1 and len(x.shape) < 3:
            return x.T
        else:
            return x

    @staticmethod
    def _return_not_none(a, b):
        """ Return the first not-None value """
        if a is None:
            return b
        else:
            return a

    @staticmethod
    def setup_array(ssm_array, y=None, skip_check_observed=False):
        """ Prepare an array of Ssm objects for other signal processing methods """
        # Convert to an array if the input is a single Ssm instance
        if hasattr(ssm_array, 'nmodel'):  # suggests that this is a single somata class object
            ssm_array = ssm_array.stack_to_array()

        # Fill the models with the same observed data if provided
        if y is not None:
            for m in range(len(ssm_array)):
                ssm_array[m].y = y

        # Verify that observed data are identical across the array
        first_ssm: StateSpaceModel = ssm_array[0]
        if not skip_check_observed:
            for ii in range(1, len(ssm_array)):
                first_ssm._check_observed_data(ssm_array[ii])

        # Array dimensions
        K = len(ssm_array)  # number of models
        Ts = [x.ntime for x in ssm_array]  # number of time points
        if len(np.unique(Ts)) == 1:
            T = Ts[0]
        else:
            T = Ts  # allow variable lengths of data in different objects

        return ssm_array, K, T

    # Kalman filtering and smoothing methods (E step)
    def kalman_filt_smooth(self, y=None, R_weights=None, return_dict=False, EM=False, skip_interp=True, seterr=None):
        """ Wrapper method for classical kalman filtering and smoothing """
        if seterr is not None:  # apply np.seterr() for parallel processes
            old_settings = np.seterr(**seterr)

        y = self.y if y is None else y
        x_t_n, P_t_n, P_t_tmin1_n, logL, x_t_t, P_t_t, K_t, x_t_tmin1, P_t_tmin1, fy_t_interp = kalman(
            F=self.F, Q=self.Q, mu0=self.mu0, S0=self.S0, G=self.G, R=self.R, y=y, R_weights=R_weights,
            skip_interp=skip_interp)

        if seterr is not None:
            # noinspection PyUnboundLocalVariable
            np.seterr(**old_settings)

        if EM:  # return minimally necessary variables for EM algorithm
            return {'x_t_n': x_t_n, 'P_t_n': P_t_n, 'P_t_tmin1_n': P_t_tmin1_n, 'logL': logL}
        else:
            if return_dict:
                return {'x_t_n': x_t_n, 'P_t_n': P_t_n, 'P_t_tmin1_n': P_t_tmin1_n, 'logL': logL,
                        'x_t_t': x_t_t, 'P_t_t': P_t_t, 'K_t': K_t,
                        'x_t_tmin1': x_t_tmin1, 'P_t_tmin1': P_t_tmin1, 'fy_t_interp': fy_t_interp}
            else:
                return x_t_n, P_t_n, P_t_tmin1_n, logL, x_t_t, P_t_t, K_t, x_t_tmin1, P_t_tmin1, fy_t_interp

    def dejong_filt_smooth(self, y=None, R_weights=None, return_dict=False, EM=False, skip_interp=True, seterr=None):
        """ Wrapper method for De Jong version kalman filtering and smoothing """
        if seterr is not None:  # apply np.seterr() for parallel processes
            old_settings = np.seterr(**seterr)

        y = self.y if y is None else y
        x_t_n, P_t_n, P_t_tmin1_n, logL, x_t_t, P_t_t, K_t, x_t_tmin1, P_t_tmin1, fy_t_interp = djkalman(
            F=self.F, Q=self.Q, mu0=self.mu0, S0=self.S0, G=self.G, R=self.R, y=y, R_weights=R_weights,
            skip_interp=skip_interp)

        if seterr is not None:
            # noinspection PyUnboundLocalVariable
            np.seterr(**old_settings)

        if EM:  # return minimally necessary variables for EM algorithm
            return {'x_t_n': x_t_n, 'P_t_n': P_t_n, 'P_t_tmin1_n': P_t_tmin1_n, 'logL': logL}
        else:
            if return_dict:
                return {'x_t_n': x_t_n, 'P_t_n': P_t_n, 'P_t_tmin1_n': P_t_tmin1_n, 'logL': logL,
                        'x_t_t': x_t_t, 'P_t_t': P_t_t, 'K_t': K_t,
                        'x_t_tmin1': x_t_tmin1, 'P_t_tmin1': P_t_tmin1, 'fy_t_interp': fy_t_interp}
            else:
                return x_t_n, P_t_n, P_t_tmin1_n, logL, x_t_t, P_t_t, K_t, x_t_tmin1, P_t_tmin1, fy_t_interp

    @staticmethod
    def par_kalman(ssm_array, y=None, method='kalman', skip_check_observed=False,
                   R_weights=None, skip_interp=True, return_dict=False):
        """
        Parallel run kalman filtering and smoothing on
        an array of StateSpaceModel objects
        """
        ssm_array, K, _ = StateSpaceModel.setup_array(ssm_array, y=y, skip_check_observed=skip_check_observed)

        if method == 'kalman':
            kalman_func = getattr(StateSpaceModel, 'kalman_filt_smooth')
        elif method == 'dejong':
            kalman_func = getattr(StateSpaceModel, 'dejong_filt_smooth')
        else:
            raise ValueError('Specified method is invalid for parallel kalman calls.')

        if R_weights is None:
            R_weights = np.tile(np.array([None]), K)
        else:
            R_weights = tuple([R_weights[m, :] for m in range(K)])

        # Run kalman filtering and smoothing on parallel processes
        n_jobs = max(cpu_count()-1, 1)
        results = Parallel(n_jobs=n_jobs)(delayed(kalman_func)(model, y, weights, False, False,
                                                               skip_interp, np.geterr())
                                          for model, weights in zip(ssm_array, R_weights))

        # Unpack results into separate variables to return
        (x_t_n_all, P_t_n_all, P_t_tmin1_n_all,
         x_t_t_all, P_t_t_all, K_t_all, x_t_tmin1_all, P_t_tmin1_all,
         logL_all, fy_t_interp_all) = tuple([[None]*K for _ in range(10)])
        for m in range(K):
            (x_t_n_all[m], P_t_n_all[m], P_t_tmin1_n_all[m],
             logL_all[m], x_t_t_all[m], P_t_t_all[m], K_t_all[m],
             x_t_tmin1_all[m], P_t_tmin1_all[m], fy_t_interp_all[m]) = results[m]

        if return_dict:
            return {'x_t_n_all': x_t_n_all, 'P_t_n_all': P_t_n_all, 'P_t_tmin1_n_all': P_t_tmin1_n_all,
                    'logL_all': logL_all, 'x_t_t_all': x_t_t_all, 'P_t_t_all': P_t_t_all, 'K_t_all': K_t_all,
                    'x_t_tmin1_all': x_t_tmin1_all, 'P_t_tmin1_all': P_t_tmin1_all, 'fy_t_interp_all': fy_t_interp_all}
        else:
            return x_t_n_all, P_t_n_all, P_t_tmin1_n_all, logL_all, \
                x_t_t_all, P_t_t_all, K_t_all, x_t_tmin1_all, P_t_tmin1_all, fy_t_interp_all

    # Parameter estimation methods (M step)
    # noinspection PyUnusedLocal
    def m_estimate(self, y=None, x_t_n=None, P_t_n=None, P_t_tmin1_n=None, h_t=None, logL=None,
                   priors=None, A=None, B=None, C=None, T=None, force_ABC=False,
                   update_param=('F', 'Q', 'mu0', 'S0', 'G', 'R'), keep_param=(),
                   return_dict=None):
        """
        Maximum likelihood or Maximum a posteriori estimation to update
        parameters. Ssm class doesn't have explicit state-equation parameter
        update rules, therefore calling the inherited _m_update_<param>
        methods in component subclasses

        Reference:
            Shumway, R. H., & Stoffer, D. S. (1982). An approach to time series
            smoothing and forecasting using the EM algorithm. Journal of time series
            analysis, 3(4), 253-264.

            Ghahramani, Z., & Hinton, G. E. (1996). Parameter estimation for linear
            dynamical systems.

            Ghahramani, Z., & Hinton, G. E. (2000). Variational learning for
            switching state-space models. Neural computation, 12(4), 831-864.

        Inputs:
        :param self: Ssm class instance
        :param y: observed data
        :param x_t_n: smoothed estimates (posterior) of state mean
        :param P_t_n: smoothed estimates (posterior) of state conditional covariance
        :param P_t_tmin1_n: smoothed estimates (posterior) of state lag1 conditional cross-covariance
        :param h_t: responsibility vector (i.e. 1/R_weights in E step)
        :param logL: unused variable, pure syntactic sugar to allow pythonic **d inputs
        :param priors: a list of dictionaries specifying priors for each component, if None -> MLE
        :param A: sum of square terms
        :param B: sum of square terms
        :param C: sum of square terms
        :param T: number of time points in sum of square terms
        :param force_ABC: flag to compute ABC regardless of state-equation parameters
        :param update_param: a tuple of strings for parameters to update
        :param keep_param: a tuple of strings for parameters to keep and not update
        :param return_dict: None -> no return, True -> return dict, False -> return tuple of variables
        """
        # Obtain boolean flags of scopes of updates
        update_FQ, update_mu0S0G = self._m_estimate_scope(update_param, keep_param)

        # Initialize parameters
        y = self.y if y is None else y
        if priors is None:
            priors = [None] * self.ncomp
        elif isinstance(priors, dict):
            priors = [priors]  # so that it can be indexed at position 0

        # Attempt to skip sums of squares computation
        if update_FQ or force_ABC:
            A = self._m_compute_ss('A', x_t_n, P_t_n=P_t_n) if A is None else A
            B = self._m_compute_ss('B', x_t_n, P_t_tmin1_n=P_t_tmin1_n) if B is None else B
            C = self._m_compute_ss('C', x_t_n, P_t_n=P_t_n) if C is None else C
            T = x_t_n.shape[1] - 1 if T is None else T

        # Initialize the responsibility vector h_t
        if h_t is None:
            h_t_length = T if T is not None else y.shape[1]
            h_t = np.ones(h_t_length, dtype=np.float64)  # default responsibility is 1 for all timepoints

        # Update parameters for each independent component -- F, Q, mu0, S0, G (component specific priors)
        if update_FQ or update_mu0S0G:
            for ii in range(self.ncomp):  # iterate through components
                current_component: StateSpaceModel = self.components[ii]  # typehint to superclass to be flexible
                start_idx = sum(self.comp_nstates[:ii])
                end_idx = sum(self.comp_nstates[:ii+1])
                if update_FQ:
                    A_tmp = A[start_idx:end_idx, start_idx:end_idx]
                    B_tmp = B[start_idx:end_idx, start_idx:end_idx]
                    C_tmp = C[start_idx:end_idx, start_idx:end_idx]
                else:
                    A_tmp, B_tmp, C_tmp = (None, None, None)

                # Call the _m_update_<param> methods specific to the component subclass
                if 'F' in update_param and 'F' not in keep_param:
                    self.F[start_idx:end_idx, start_idx:end_idx] = \
                        current_component._m_update_f(A=A_tmp, B=B_tmp, priors=priors[ii])

                if 'Q' in update_param and 'Q' not in keep_param:
                    self.Q[start_idx:end_idx, start_idx:end_idx] = \
                        current_component._m_update_q(A=A_tmp, B=B_tmp, C=C_tmp, T=T,
                                                      F=self.F[start_idx:end_idx, start_idx:end_idx],
                                                      priors=priors[ii])

                if 'mu0' in update_param and 'mu0' not in keep_param:
                    self.mu0[start_idx:end_idx, 0] = \
                        current_component._m_update_mu0(x_0_n=x_t_n[start_idx:end_idx, 0])[:, 0]

                if 'S0' in update_param and 'S0' not in keep_param:
                    self.S0[start_idx:end_idx, start_idx:end_idx] = \
                        current_component._m_update_S0(x_0_n=x_t_n[start_idx:end_idx, 0],
                                                       P_0_n=P_t_n[start_idx:end_idx, start_idx:end_idx, 0],
                                                       mu0=self.mu0[start_idx:end_idx, 0][:, None])

                if 'G' in update_param and 'G' not in keep_param:
                    comp_G = current_component._m_update_g(y=y, x_t_n=x_t_n[start_idx:end_idx, :],
                                                           P_t_n=P_t_n[start_idx:end_idx, start_idx:end_idx, :],
                                                           h_t=h_t)
                    self.G[:, start_idx:end_idx] = comp_G if comp_G is not None else self.G[:, start_idx:end_idx]

        # Update observation noise covariance -- R
        if 'R' in update_param and 'R' not in keep_param:
            R_ss = self._m_update_r(y=y, x_t_n=x_t_n, P_t_n=P_t_n, h_t=h_t, G=self.G, priors=priors[0])
        else:
            R_ss = self._m_update_r(y=y, x_t_n=x_t_n, P_t_n=P_t_n, h_t=h_t, G=self.G, keep_R=True)

        if return_dict is None:
            pass
        elif return_dict:
            return {'F': self.F, 'Q': self.Q, 'mu0': self.mu0, 'S0': self.S0, 'G': self.G, 'R': self.R,
                    'R_ss': R_ss, 'A': A, 'B': B, 'C': C}
        else:
            return self.F, self.Q, self.mu0, self.S0, self.G, self.R, R_ss, A, B, C

    def update_comp_param(self):
        """ Update component specific parameters, override in subclasses """
        pass

    def e_step(self, y=None, logL_list=None, **kwargs):
        """ Exposed E step method for run_em() """
        e_results: dict = self.dejong_filt_smooth(y=y, EM=True, **kwargs)
        e_logL = e_results['logL'].sum()
        stop_var = e_logL - logL_list[-1]  # logL should monotonically increase
        logL_list.append(e_logL)
        return e_results, stop_var

    def m_step(self, y=None, **kwargs):
        """ Exposed M step method for run_em() """
        self.m_estimate(y=y, **kwargs)
        return

    def initialize_priors(self):
        """ Initialize priors for each component in the object """
        assert self.components is not None, 'Cannot initialize priors when components is None.'
        components_prefill = self.fill_components()
        priors = [x.initialize_priors() for x in self.components]
        self.unfill_components(components_prefill)
        return priors

    def _initialize_priors_recursive_list(self, prior_value):
        """ Expand a prior value to the length of components for the recursive case """
        if prior_value is None:
            prior_value = [None] * self.ncomp
        elif isinstance(prior_value, str) or not isinstance(prior_value, Iterable) or len(prior_value) == 1:
            prior_value = [prior_value] * self.ncomp
        else:
            assert len(prior_value) == self.ncomp, 'Different lengths of specified prior value and components.'
        return prior_value

    def _m_update_r(self, y=None, x_t_n=None, P_t_n=None, h_t=None, G=None, R_ss=None, T=None, priors=None,
                    keep_R=False):
        """ Update observation noise covariance """
        if R_ss is None:
            assert len(h_t) == y.shape[1], 'Different lengths of h_t and y. Cannot proceed with _m_update_r().'
            R_ss = (h_t * (y - G @ x_t_n[:, 1:])) @ (y - G @ x_t_n[:, 1:]).T + \
                G @ (h_t * P_t_n[:, :, 1:]).sum(axis=2) @ G.T

        if not keep_R:  # proceed to update self.R
            T = h_t.sum() if T is None else T

            if self._m_update_if_mle('R_sigma2', priors):
                # MLE
                self.R = R_ss / T
            else:
                # MAP with inverse gamma prior
                R_init = priors['R_sigma2']
                R_hp = priors['R_hyperparameter'] if 'R_hyperparameter' in priors else 0.1
                alpha = T * R_hp / 2  # scales with data length T according to the hyperparameter
                beta = R_init * (alpha + 1)  # setting the mode of inverse gamma prior to be R_init
                self.R = (beta + R_ss/2) / (alpha + T/2 + 1)  # new R is the mode of inverse gamma posterior

        return R_ss

    @staticmethod
    def _m_update_if_mle(params, priors):
        """ Check whether doing MLE during _m_update_<param> methods """
        if priors is None:
            return True

        if isinstance(params, str) or not isinstance(params, Iterable):
            params = [params]

        MLE_flag = np.zeros(len(params), dtype=bool)
        for param, idx in zip(params, range(len(params))):
            if param not in priors:
                MLE_flag[idx] = True
            elif isinstance(priors[param], str):
                if priors[param].upper() == 'MLE':
                    MLE_flag[idx] = True

        return MLE_flag.all()

    # noinspection PyUnusedLocal
    @staticmethod
    def _m_update_f(A=None, B=None, priors=None):
        pass

    # noinspection PyUnusedLocal
    @staticmethod
    def _m_update_q(A=None, B=None, C=None, T=None, F=None, priors=None):
        pass

    # noinspection PyUnusedLocal
    @staticmethod
    def _m_update_q_src(Q_ss=None, T=None, Q_basis=None, nsource=None, nstate=None, priors=None):
        pass

    @staticmethod
    def _m_update_mu0(x_0_n=None):
        """ Update initial state mean -- mu0 """
        mu0 = x_0_n[:, None].copy()
        return mu0

    @staticmethod
    def _m_update_mu0_src(x_0_n=None, nstate=None):
        """
        Update initial state mean mu0 in dynamic
        source localization
        """
        assert x_0_n.shape[0] % nstate == 0, 'Input x_t_n is not a multiple of nstate.'
        mu0 = x_0_n[:, None].clone()
        return mu0

    @staticmethod
    def _m_update_S0(x_0_n=None, P_0_n=None, mu0=None):
        """ Update initial state covariance -- S0 """
        S0 = P_0_n + x_0_n[:, None] @ x_0_n[:, None].T \
            - x_0_n[:, None] @ mu0.T - mu0 @ x_0_n[:, None].T + mu0 @ mu0.T
        return S0

    @staticmethod
    def _m_update_g(y=None, x_t_n=None, P_t_n=None, h_t=None, C=None, D=None):
        """ Update observation matrix -- G """
        if C is None:
            assert len(h_t) == y.shape[1], 'Different lengths of h_t and y. Cannot proceed with _m_update_g().'
            C = (h_t * P_t_n[:, :, 1:]).sum(axis=2) + (h_t * x_t_n[:, 1:]) @ x_t_n[:, 1:].T

        if D is None:
            assert len(h_t) == y.shape[1], 'Different lengths of h_t and y. Cannot proceed with _m_update_g().'
            D = (h_t * y) @ x_t_n[:, 1:].T

        approach = 'svd' if C.shape[0] >= 5 else 'gaussian'
        G = D @ inverse(C, approach=approach)
        return G

    @staticmethod
    def _m_estimate_scope(update_param=('F', 'Q', 'mu0', 'S0', 'G', 'R'), keep_param=()):
        """ Sort out the scopes of updates in m_estimate() """
        if 'F' not in update_param and 'Q' not in update_param:
            update_FQ = False
        elif 'F' in keep_param and 'Q' in keep_param:
            update_FQ = False
        else:
            update_FQ = True

        if 'mu0' not in update_param and 'S0' not in update_param and 'G' not in update_param:
            update_mu0S0G = False
        elif 'mu0' in keep_param and 'S0' in keep_param and 'G' in keep_param:
            update_mu0S0G = False
        else:
            update_mu0S0G = True

        return update_FQ, update_mu0S0G

    @staticmethod
    def _m_compute_ss(ss_term, x_t_n, P_t_n=None, P_t_tmin1_n=None):
        """
        Compute the sum of squares terms for m_estimate().
        Definitions of A,B,C follow the notations in equations (9,10,11)
        of S&S 1982, and these terms also correspond to maximizing <H>_Q
        in G&H 2000.

        Reference:
            Shumway, R. H., & Stoffer, D. S. (1982). An approach to time series
            smoothing and forecasting using the EM algorithm. Journal of time series
            analysis, 3(4), 253-264.

            Ghahramani, Z., & Hinton, G. E. (2000). Variational learning for
            switching state-space models. Neural computation, 12(4), 831-864.
        """
        if ss_term == 'A':
            return P_t_n[:, :, :-1].sum(axis=2) + x_t_n[:, :-1] @ x_t_n[:, :-1].T
        elif ss_term == 'B':
            return P_t_tmin1_n[:, :, 1:].sum(axis=2) + x_t_n[:, 1:] @ x_t_n[:, :-1].T
        elif ss_term == 'C':
            return P_t_n[:, :, 1:].sum(axis=2) + x_t_n[:, 1:] @ x_t_n[:, 1:].T
        else:
            raise ValueError('Invalid ss_term is specified.')
