#!/usr/bin/env python
"""
In bayesloop, each new data study is handled by an instance of a ``Study``-class. In this way, all data, the inference
results and the appropriate post-processing routines are stored in one object that can be accessed conveniently or
stored in a file. Apart from the basic ``Study`` class, there exist a number of specialized classes that extend the
basic fit method, for example to infer the full distribution of hyper-parameters or to apply model selection to on-line
data streams.
"""

from __future__ import division, print_function
import string
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from scipy.special import factorial, logsumexp
from scipy.special import beta as beta_func
import sympy.abc as abc
from sympy import symbols
from sympy import lambdify
from sympy.stats import density
import sympy.stats
from copy import copy, deepcopy
from collections import OrderedDict
from collections.abc import Iterable
from inspect import getargspec
from tqdm.auto import tqdm
from .helper import assignNestedItem, recursiveIndex, flatten, createColormap, oint, cint, freeSymbols
from .preprocessing import movingWindow
from .observationModels import ObservationModel
from .transitionModels import TransitionModel, CombinedTransitionModel, SerialTransitionModel
from .exceptions import ConfigurationError, PostProcessingError
from .parser import Parser


class Study(object):
    """
    Fits with fixed hyper-parameters and hyper-parameter optimization. This class implements a
    forward-backward-algorithm for analyzing time series data using hierarchical models. For efficient computation,
    all parameter distributions are discretized on a parameter grid.
    """
    def __init__(self, silent=False):
        self.observationModel = None
        self.transitionModel = None

        self.gridSize = []
        self.boundaries = []
        self.marginalGrid = []
        self.grid = []
        self.latticeConstant = []

        self.rawData = np.array([])
        self.formattedData = np.array([])
        self.rawTimestamps = None
        self.formattedTimestamps = None

        self.posteriorSequence = []
        self.posteriorMeanValues = []
        self.logEvidence = 0
        self.localEvidence = []

        self.selectedHyperParameters = []
        self.fitWarningCounter = 0

        if not silent:
            print('+ Created new study.')

    @property
    def log10Evidence(self):
        return self.logEvidence/np.log(10)

    def loadExampleData(self, silent=False):
        """
        Loads UK coal mining disaster data.

        Args:
            silent(bool): If set to True, no output is generated by this method.
        """
        self.rawData = np.array([5, 4, 1, 0, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6, 3, 3, 5, 4, 5, 3, 1, 4,
                                 4, 1, 5, 5, 3, 4, 2, 5, 2, 2, 3, 4, 2, 1, 3, 2, 2, 1, 1, 1, 1, 3, 0,
                                 0, 1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1, 0, 1, 0, 1, 0, 0,
                                 0, 2, 1, 0, 0, 0, 1, 1, 0, 2, 3, 3, 1, 1, 2, 1, 1, 1, 1, 2, 3, 3, 0,
                                 0, 0, 1, 4, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0])

        self.rawTimestamps = np.arange(1852, 1962)

        if not silent:
            print('+ Successfully imported example data.')

    def loadData(self, array, timestamps=None, silent=False):
        """
        Loads Numpy array as data.

        Args:
            array(ndarray): Numpy array containing time series data
            timestamps(ndarray): Array of timestamps (same length as data array)
            silent(bool): If set to True, no output is generated by this method.
        """
        if isinstance(array, np.ndarray):
            self.rawData = array
        elif isinstance(array, list):
            if not silent:
                print('! WARNING: Data supplied as list, not as Numpy array. Converting list to Numpy array '
                      '(dtype=float).')
            self.rawData = np.array(array, dtype=float)
        else:
            raise ConfigurationError('Data type not supported. Please provide data as Numpy array.')

        if timestamps is not None:  # load custom timestamps
            if len(timestamps) == len(array):
                self.rawTimestamps = np.array(timestamps)
            else:
                if not silent:
                    print('! WARNING: Number of timestamps does not match number of data points. Omitting timestamps.')
        else:  # set default timestamps (integer range)
            self.rawTimestamps = np.arange(len(self.rawData))

        if not silent:
            print('+ Successfully imported array.')

    def load(self, array, timestamps=None, silent=False):
        """
        See :meth:`.Study.loadData`.
        """
        self.loadData(array, timestamps=timestamps, silent=silent)

    def setObservationModel(self, L, silent=False):
        """
        Sets observation model (likelihood function) for analysis and creates parameter grid for inference routine.

        Args:
            L: Observation model class (see observationModels.py)
            silent(bool): If set to True, no output is generated by this method.
        """
        self.observationModel = L

        # prepare parameter grid
        self.marginalGrid = []
        self.gridSize = []
        self.boundaries = []
        self.latticeConstant = []
        for v, n in zip(self.observationModel.parameterValues, self.observationModel.parameterNames):
            if v is None:  # if user has not specified parameter values, we try to estimate them
                try:
                    v = self.observationModel.estimateParameterValues(n, self.rawData)
                    print('+ Estimated parameter interval for "{}": [{}, {}] ({} values).'
                          .format(n, v[0], v[-1], len(v)))
                except:
                    raise ConfigurationError('Could not estimate parameter values for "{}".'.format(n))

            v = np.array(v, dtype=float)  # inference algorithm needs floats

            self.marginalGrid.append(v)
            self.gridSize.append(len(v))
            self.boundaries.append([v[0], v[-1]])

            # check if parameter values are equally spaced
            if np.any(np.abs(np.diff(np.diff(v))) > 10 ** -10):
                print('! WARNING: Supplied parameter values for "{}" are not equally spaced. Assuming categorical '
                      'parameter.'.format(n))
                self.latticeConstant.append(1.)
            else:  # equally spaced (regular grid)
                self.latticeConstant.append(np.abs(v[0] - v[1]))

        # create grid
        self.grid = [m for m in np.meshgrid(*self.marginalGrid, indexing='ij')]

        # if observation model is updated, transition model must know the new lattice constant
        if self.transitionModel is not None:
            self.transitionModel.latticeConstant = self.latticeConstant

        if not silent:
            print('+ Observation model: {}. Parameter(s): {}'.format(L, L.parameterNames))

    def setOM(self, L, silent=False):
        """
        See :meth:`.Study.setObservationModel`.
        """
        self.setObservationModel(L, silent=silent)

    def _computePrior(self, silent=False):
        """
        Computes discrete prior probabilities (densities) for the parameters of the observation model. The custom prior
        distribution may be passed as a Numpy array that has the same shape as the parameter grid, as a (lambda)
        function or as a (list of) SymPy random variable(s).

        Args:
            silent(bool): If set to True, no output is generated by this method.

        Returns:
            ndarray: Prior probability (density) values with the same size as the parameter grid
        """
        prior = self.observationModel.prior

        # check if no prior is specified
        if prior is None:
            prior = np.ones(self.gridSize)  # flat (improper) prior
            prior /= np.sum(prior)  # re-normalize (now equal to uniform prior)
            prior /= np.prod(self.latticeConstant)  # convert to prob. density
            if not silent:
                print('    + Set uniform prior with parameter boundaries.')
            return prior

        # check whether correctly shaped numpy array is provided
        if isinstance(prior, np.ndarray):
            if np.all(prior.shape == self.grid[0].shape):
                norm = np.sum(prior)
                if norm != 1.:
                    prior /= norm
                    prior /= np.prod(self.latticeConstant)
                    if not silent:
                        print('    + Set prior (numpy array). Values have been re-normalized.')
                else:
                    if not silent:
                        print('    + Set prior (numpy array).')
                return prior
            else:
                raise ConfigurationError('Prior array does not match parameter grid size.')

        # check whether function is provided
        if hasattr(prior, '__call__'):
            values = prior(*self.grid)*np.ones(self.gridSize)  # last factor for cases like lambda x: 1.
            norm = np.sum(values)
            if norm != 1.:
                values /= norm
                values /= np.prod(self.latticeConstant)
                if not silent:
                    print('    + Set prior (function): {}. Values have been re-normalized.'.format(prior.__name__))
            else:
                if not silent:
                    print('    + Set prior (function): {}'.format(prior.__name__))
            return values

        # check whether single random variable is provided
        if type(prior) is sympy.stats.rv.RandomSymbol:
            prior = [prior]

        # check if list/tuple is provided
        if isinstance(prior, (list, tuple)) and not isinstance(prior, str):
            if len(prior) != len(self.observationModel.parameterNames):
                raise ConfigurationError('Observation model contains {} parameters, but {} priors were provided.'
                                         .format(len(self.observationModel.parameterNames), len(prior)))

            pdf = 1
            if len(prior) > 1:
                x = symbols(' '.join(list(string.ascii_lowercase)[:len(prior)]))
            else:
                x = [abc.x]

            for i, rv in enumerate(prior):
                if type(rv) is not sympy.stats.rv.RandomSymbol:
                    raise ConfigurationError('Only lambda functions or SymPy random variables can be used as a prior.')
                if len(freeSymbols(rv)) > 0:
                    raise ConfigurationError('Prior distribution must not contain free parameters.')

                # multiply total pdf with density for current parameter
                pdf = pdf*density(rv)(x[i])

            # set density as lambda function
            if not silent:
                print('    + Set prior (sympy): {}'.format(pdf))
            return lambdify(x, pdf, modules=['numpy', {'factorial': factorial, 'beta': beta_func}])(*self.grid)

    def setTransitionModel(self, T, silent=False):
        """
        Set transition model which describes the parameter dynamics.

        Args:
            T: Transition model class (see transitionModels.py)
            silent(bool): If true, no output is printed by this method
        """
        # check if model is a break-point and raise error if so
        if str(T) == 'Break-point':
            raise ConfigurationError('The "BreakPoint" transition model can only be used with the '
                                     '"SerialTransitionModel" class.')

        self.transitionModel = T
        self.transitionModel.study = self
        self.transitionModel.latticeConstant = self.latticeConstant
        if not silent:
            print('+ Transition model: {}. Hyper-Parameter(s): {}'
                  .format(T, self._unpackAllHyperParameters(values=False)))

    def setTM(self, T, silent=False):
        """
        See :meth:`.Study.setTransitionModel`.
        """
        self.setTransitionModel(T, silent=silent)

    def set(self, *args, **kwargs):
        """
        Set observation model or transition model, or both. See :meth:`.Study.setTransitionModel` and
        :meth:`.Study.setObservationModel`.

        Args:
            args: Sequence of Observation model instance and Transition model instance, or just one of those two types
            silent(bool): If true, no output is printed by this method
        """
        # Check for unknown keyword-arguments
        for key in kwargs.keys():
            if key not in ['silent']:
                raise TypeError("set() got an unexpected keyword argument '{}'".format(key))

        # For Python 2 compatibility (no argument with default value after *args)
        silent = kwargs.pop('silent', False)

        om = False
        tm = False

        for model in args:
            if ObservationModel in model.__class__.__bases__:
                if om:
                    raise ConfigurationError('More than one observation model supplied.')
                om = True
                self.setObservationModel(model, silent=silent)

            elif TransitionModel in model.__class__.__bases__:
                if tm:
                    raise ConfigurationError('More than one transition model supplied.')

                tm = True
                self.setTransitionModel(model, silent=silent)

            else:
                raise ConfigurationError('Expected observation model or transition model instance as first argument.')

    def fit(self, forwardOnly=False, evidenceOnly=False, silent=False):
        """
        Computes the sequence of posterior distributions and evidence for each time step. Evidence is also computed for
        the complete data set.

        Args:
            forwardOnly(bool): If set to True, the fitting process is terminated after the forward pass. The resulting
                posterior distributions are so-called "filtering distributions" which - at each time step -
                only incorporate the information of past data points. This option thus emulates an online
                analysis.
            evidenceOnly(bool): If set to True, only forward pass is run and evidence is calculated. In contrast to the
                forwardOnly option, no posterior mean values are computed and no posterior distributions are stored.
            silent(bool): If set to True, no output is generated by the fitting method.
        """
        self._checkConsistency()

        if not silent:
            print('+ Started new fit:')

        self.formattedData = movingWindow(self.rawData, self.observationModel.segmentLength)
        self.formattedTimestamps = self.rawTimestamps[self.observationModel.segmentLength-1:]
        if not silent:
            print('    + Formatted data.')

        # initialize array for posterior distributions
        if not evidenceOnly:
            self.posteriorSequence = np.empty([len(self.formattedData)]+self.gridSize)

        # initialize array for computed evidence (marginal likelihood)
        self.logEvidence = 0
        self.localEvidence = np.empty(len(self.formattedData))

        # set prior distribution for forward-pass
        alpha = self._computePrior(silent=silent)

        # show progressbar if silent=False
        if not silent:
            enum = tqdm(np.arange(0, len(self.formattedData)), total=len(self.formattedData))
        else:
            enum = np.arange(0, len(self.formattedData))

        # forward pass
        for i in enum:

            # compute likelihood
            likelihood = self.observationModel.processedPdf(self.grid, self.formattedData[i])

            # force dtype float on likelihood (in case it is of dtype object)
            if likelihood.dtype == object:
                likelihood = likelihood.astype(float)

            # update alpha based on likelihood
            alpha *= likelihood

            # normalization constant of alpha is used to compute evidence
            norm = np.sum(alpha)

            # normalize alpha (for numerical stability)
            if norm > 0.:
                alpha /= norm
            else:
                # if all probability values are zero, normalization is not possible
                if self.fitWarningCounter < 5:
                    print('    ! WARNING: Forward pass distribution contains only zeros, check parameter boundaries!')
                    print('      Stopping inference process. Setting model evidence to zero.')
                elif self.fitWarningCounter == 5:
                    print('    ! WARNING: Will omit further warnings about parameter boundaries.')

                self.fitWarningCounter += 1
                self.logEvidence = -np.inf
                return

            # update log-evidence and compute local evidence
            self.logEvidence += np.log(norm)
            self.localEvidence[i] = norm * np.prod(self.latticeConstant)  # integration yields evidence, not only sum

            # alphas are stored as preliminary posterior distributions
            if not evidenceOnly:
                self.posteriorSequence[i] = alpha

            # compute alpha for next iteration
            alpha = self.transitionModel.computeForwardPrior(alpha, self.formattedTimestamps[i])

        # remove progressbar correctly
        if not silent:
            enum.close()

        self.logEvidence += np.log(np.prod(self.latticeConstant))  # integration yields evidence, not only sum
        if not silent:
            print('    + Finished forward pass.')
            print('    + Log10-evidence: {:.5f}'.format(self.logEvidence / np.log(10)))

        if not (forwardOnly or evidenceOnly):
            # initialize beta and normalize by grid size (for numerical stability only)
            beta = np.ones(self.gridSize)
            beta /= np.sum(beta)

            # show progressbar if silent=False
            if not silent:
                enum = tqdm(np.arange(0, len(self.formattedData))[::-1], total=len(self.formattedData))
            else:
                enum = np.arange(0, len(self.formattedData))[::-1]

            # backward pass
            for i in enum:
                # posterior ~ alpha*beta
                self.posteriorSequence[i] *= beta  # alpha*beta

                # normalize posterior wrt the parameters
                norm = np.sum(self.posteriorSequence[i])
                if norm > 0.:
                    self.posteriorSequence[i] /= np.sum(self.posteriorSequence[i])
                else:
                    # if all posterior probabilities are zero, normalization is not possible
                    if self.fitWarningCounter < 5:
                        print('    ! WARNING: Posterior distribution contains only zeros, check parameter boundaries!')
                        print('      Stopping inference process. Setting model evidence to zero.')
                    elif self.fitWarningCounter == 5:
                        print('    ! WARNING: Will omit further warnings about parameter boundaries.')

                    self.fitWarningCounter += 1
                    self.logEvidence = -np.inf
                    return

                # re-compute likelihood
                likelihood = self.observationModel.processedPdf(self.grid, self.formattedData[i])

                # force dtype float on likelihood (in case it is of dtype object)
                if likelihood.dtype == object:
                    likelihood = likelihood.astype(float)

                # compute local evidence
                with np.errstate(invalid='ignore'):
                    self.localEvidence[i] = 1./(np.sum(self.posteriorSequence[i]/likelihood) *
                                                np.prod(self.latticeConstant))  # integration, not only sum

                # compute beta for next iteration
                beta = self.transitionModel.computeBackwardPrior(beta*likelihood, self.formattedTimestamps[i])

                # normalize beta (for numerical stability)
                beta /= np.sum(beta)

            if not silent:
                enum.close()  # remove progressbar correctly
                print('    + Finished backward pass.')

        # posterior mean values do not need to be computed for evidence
        if evidenceOnly:
            self.posteriorMeanValues = []
        else:
            self.posteriorMeanValues = np.empty([len(self.grid), len(self.posteriorSequence)])

            for i in range(len(self.grid)):
                self.posteriorMeanValues[i] = np.array([np.sum(p*self.grid[i]) for p in self.posteriorSequence])

            if not silent:
                print('    + Computed mean parameter values.')

    def optimize(self, parameterList=[], forwardOnly=False, **kwargs):
        """
        Uses the COBYLA minimization algorithm from SciPy to perform a maximization of the log-evidence with respect
        to all hyper-parameters (the parameters of the transition model) of a time seris model. The starting values
        are the values set by the user when defining the transition model.

        For the optimization, only the log-evidence is computed and no parameter distributions are stored. When a local
        maximum is found, the parameter distribution is computed based on the optimal values for the hyper-parameters.

        Args:
            parameterList(list): List of hyper-parameter names to optimize. For nested transition models with multiple,
                identical hyper-parameter names, the sub-model index can be provided. By default, all hyper-parameters
                are optimized.
            forwardOnly(bool): If set to True, the fitting process is terminated after the forward pass. The resulting
                posterior distributions are so-called "filtering distributions" which - at each time step -
                only incorporate the information of past data points.
            **kwargs - All other keyword parameters are passed to the 'minimize' routine of scipy.optimize.
        """
        # set list of parameters to optimize
        if isinstance(parameterList, str):  # in case only a single parameter name is provided as a string
            self.selectedHyperParameters = [parameterList]
        else:
            self.selectedHyperParameters = parameterList

        print('+ Starting optimization...')
        self._checkConsistency()

        if self.selectedHyperParameters:
            print('  --> Parameter(s) to optimize:', self.selectedHyperParameters)
        else:
            print('  --> All model parameters are optimized (except change/break-points).')

            # load all hyper-parameter names (but remove break- and change-points)
            allHyperParameters = list(flatten(self._unpackHyperParameters(self.transitionModel)))
            points = list(flatten(self._unpackChangepointNames(self.transitionModel))) + \
                     list(flatten(self._unpackBreakpointNames(self.transitionModel)))
            self.selectedHyperParameters = [x for x in allHyperParameters if x not in points]

        # create parameter list to set start values for optimization
        x0 = self._unpackSelectedHyperParameters()

        # check if valid parameter names were entered
        if len(x0) == 0:
            # reset list of parameters to optimize, so that unpacking and setting hyper-parameters works as expected
            self.selectedHyperParameters = []
            raise ConfigurationError('No parameters to optimize. Check parameter names.')

        # perform optimization (maximization of log-evidence)
        result = minimize(self._optimizationStep, x0, method='COBYLA', **kwargs)

        print('+ Finished optimization.')

        # set optimal hyperparameters in transition model
        self._setSelectedHyperParameters(result.x)

        # run analysis with optimal parameter values
        self.fit(forwardOnly=forwardOnly)

        # reset list of parameters to optimize, so that unpacking and setting hyper-parameters works as expected
        self.selectedHyperParameters = []

    def _optimizationStep(self, x):
        """
        Wrapper for the fit method to use it in conjunction with scipy.optimize.minimize.

        Args:
            x(list): unpacked list of current hyper-parameter values
        """
        # set new hyperparameters in transition model
        self._setSelectedHyperParameters(x)

        # compute log-evidence
        self.fit(evidenceOnly=True, silent=True)

        print('    + Log10-evidence: {:.5f}'.format(self.logEvidence / np.log(10)), '- Parameter values:', x)

        # return negative log-evidence (is minimized to maximize evidence)
        return -self.logEvidence

    def simulate(self, x, t=None, density=False):
        """
        Computes the probability (density) for a set of observations, based on the inferred parameter distributions of a
        given time step, or based on the time-averaged parameter distributions. It can be used to compute the expected
        distribution of the observed data, taking into account the uncertainty in the parameters (as well as
        hyper-parameters for Hyper-Studies).

        Args:
            x: array of observation values
            t: Time step/stamp for which the parameter distribution is evaluated
            density: If true, probability density is computed; if false, probability value is computed

        Returns:
            ndarray: probability (density) values corresponding to observation values
        """
        if self.observationModel.segmentLength > 1:
            raise NotImplementedError('Method "simulate" is only available for observation models with '
                                      'segment length 1.')

        # if no time is provided, use time-averaged posterior distribution
        if t is None:
            post = np.sum(self.posteriorSequence, axis=0) / len(self.posteriorSequence)
        else:
            # check if supplied time stamp exists
            if t not in self.formattedTimestamps:
                raise PostProcessingError('Supplied time ({}) does not exist in data or is out of range.'.format(t))
            timeIndex = list(self.formattedTimestamps).index(t)  # to select corresponding posterior distribution
            post = self.posteriorSequence[timeIndex]

        # compute distribution of observations/data at the given points x
        prob = np.array([np.sum(self.observationModel.pdf(self.grid, [xi]) * post) for xi in x])

        if not density:
            prob /= np.sum(prob)

        return prob

    def eval(self, query, t=None, silent=False):
        """
        Convenience method to evaluate arithmetic operations on (hyper-)parameters. See :meth:`parser.Parser` for more
        information. Note: This method is slow for evaluating a lot of queries subsequently, as it initializes a new
        Parser instance for each query. Use a dedicated Parser instance in that case (i.e.: P = Parser(S); P('query')).

        Args:
            query(str): (In-)equality to compute probability for, or just an arithmetic operation on parameters (the
                distribution will be returned in the latter case).
            t: Time step/stamp to evaluate all parameters at
            silent(bool): If true, no output is generated by this method.

        Returns:
            Probability of queried (in-)equality, or values and corresponding probability values of derived
            distribution.
        """
        P = Parser(self)
        return P(query, t=t, silent=silent)

    def _unpackHyperParameters(self, transitionModel, values=False):
        """
        Returns list of all hyper-parameters (names or values), nested as the transition model.

        Args:
            transitionModel: An instance of a transition model
            values: By default, parameter names are returned; if set to True, parameter values are returned

        Returns:
            list: hyper-parameters (names or values)
        """
        paramList = []
        # recursion step for sub-models
        if hasattr(transitionModel, 'models'):
            for m in transitionModel.models:
                paramList.append(self._unpackHyperParameters(m, values=values))

        # extend hyper-parameter based on current (sub-)model
        if hasattr(transitionModel, 'hyperParameterNames'):
            if values:
                paramList.extend(transitionModel.hyperParameterValues)
            else:
                paramList.extend(transitionModel.hyperParameterNames)

        return paramList

    def _unpackAllHyperParameters(self, values=True):
        """
        Returns a flattened list of all hyper-parameter values of the current transition model.

        Returns:
            list: all hyper-parameter values of the current transition model
        """
        return list(flatten(self._unpackHyperParameters(self.transitionModel, values=values)))

    def _unpackSelectedHyperParameters(self):
        """
        The parameters of a transition model can be split between several sub-models (using CombinedTransitionModel or
        SerialTransitionModel) and can be lists of values (multiple standard deviations in GaussianRandomWalk). This
        function unpacks the hyper-parameters, resulting in a single list of values that can be fed to the optimization
        step routine. Note that only the hyper-parameters that are noted (by name) in the attribute
        selectedHyperParameters are regarded.

        Returns:
            list: currently selected hyper-parameter values if successful, 0 otherwise
        """
        # if no hyper-parameters are selected, choose all
        if not self.selectedHyperParameters:
            return self._unpackAllHyperParameters()

        # if self.selectedHyperParameters is not empty
        nameTree = self._unpackHyperParameters(self.transitionModel)
        valueTree = self._unpackHyperParameters(self.transitionModel, values=True)
        output = []

        # loop over selected hyper-parameters
        for name in self.selectedHyperParameters:
            iFound = recursiveIndex(nameTree, name)  # choose first hit
            if len(iFound) == 0:
                raise ConfigurationError('Could not find any hyper-parameter named {}.'.format(name))

            value = valueTree[:]
            for i in iFound:
                value = value[i]

            output.append(value)

            # remove occurrence from nameTree (if name is listed twice, use second occurrence...)
            assignNestedItem(nameTree, iFound, ' ')

        # return selected values of hyper-parameters
        return output

    def _setAllHyperParameters(self, x):
        """
        Sets all current hyper-parameters, based on a flattened list of parameter values.

        Args:
            x(list): list of values (e.g. from _unpackSelectedHyperParameters)
        """
        paramList = list(x[:])  # make copy of parameter list

        nameTree = self._unpackHyperParameters(self.transitionModel)
        namesFlat = list(flatten(self._unpackHyperParameters(self.transitionModel)))

        for name in namesFlat:
            index = recursiveIndex(nameTree, name)

            # get correct sub-model
            model = self.transitionModel
            for i in index[:-1]:
                model = model.models[i]

            model.hyperParameterValues[model.hyperParameterNames.index(name)] = paramList[0]
            paramList.pop(0)

            # remove occurrence from nameTree (if name is listed twice, use second occurrence...)
            assignNestedItem(nameTree, index, ' ')

    def _setSelectedHyperParameters(self, x):
        """
        The parameters of a transition model can be split between several sub-models (using CombinedTransitionModel or
        SerialTransitionModel) and can be lists of values (multiple standard deviations in GaussianRandomWalk). This
        function takes a list of values and sets the corresponding variables in the transition model instance. Note that
        only the hyper-parameters that are noted (by name) in the attribute selectedHyperParameters are regarded.

        Args:
            x(list): list of values (e.g. from _unpackSelectedHyperParameters)

        Returns:
            int: 1, if successful, 0 otherwise
        """
        # if no hyper-parameters are selected, choose all
        if not self.selectedHyperParameters:
            self._setAllHyperParameters(x)
            return 1

        paramList = list(x[:])  # make copy of parameter list
        nameTree = self._unpackHyperParameters(self.transitionModel)

        # loop over selected hyper-parameters
        for name in self.selectedHyperParameters:
            iFound = recursiveIndex(nameTree, name)  # choose first hit
            if len(iFound) == 0:
                raise ConfigurationError('Could not find any hyper-parameter named {}.'.format(name))

            # get correct sub-model
            model = self.transitionModel
            for i in iFound[:-1]:
                model = model.models[i]

            model.hyperParameterValues[model.hyperParameterNames.index(name)] = paramList[0]
            paramList.pop(0)

            # remove occurrence from nameTree (if name is listed twice, use second occurrence...)
            assignNestedItem(nameTree, iFound, ' ')
        return 1

    def _unpackChangepointNames(self, transitionModel):
        """
        Returns list of all hyper-parameter names that are associated with change-points, nested like the transition
        model.

        Returns:
            list: all hyper-parameter names that are associated with change-points
        """
        paramList = []
        # recursion step for sub-models
        if hasattr(transitionModel, 'models'):
            for m in transitionModel.models:
                paramList.append(self._unpackChangepointNames(m))

        # extend hyper-parameter based on current (sub-)model
        if hasattr(transitionModel, 'hyperParameterNames'):
            if str(transitionModel) == 'Change-point':
                paramList.extend(transitionModel.hyperParameterNames)

        return paramList

    def _unpackBreakpointNames(self, transitionModel):
        """
        Returns list of all hyper-parameter names that are associated with break-points, nested like the transition
        model.

        Returns:
            list: all hyper-parameter names that are associated with break-points
        """
        paramList = []
        # recursion step for sub-models
        if hasattr(transitionModel, 'models'):
            for m in transitionModel.models:
                paramList.append(self._unpackBreakpointNames(m))

        # extend hyper-parameter based on current (sub-)model
        if hasattr(transitionModel, 'hyperParameterNames'):
            if str(transitionModel) == 'Serial transition model':
                paramList.extend(transitionModel.hyperParameterNames)

        return paramList

    def _getHyperParameterIndex(self, transitionModel, name):
        """
        Helper function that returns the index at which a hyper-parameter is found in the flattened list of
        hyper-parameter names.

        Args:
            transitionModel: transition model instance in which to search
            name(str): Name of a hyper-parameter. If the name occurs multiple times, the index of the submodel can be
                supplied (starting at 1 for the first submodel).

        Returns:
            int: index of the hyper-parameter
        """
        # no index provided: choose first occurrence and determine axis of hyper-parameter on grid of
        # hyper-parameter values
        hpn = list(flatten(self._unpackHyperParameters(transitionModel, values=False)))
        if name in hpn:
            paramIndex = hpn.index(name)
        else:
            raise PostProcessingError('Could not find any hyper-parameter with name: {}.'.format(name))

        return paramIndex

    def getHyperParameterValue(self, name):
        """
        Returns the currently set value of a hyper-parameter. Note: The returned value is NOT an inferred value, but
        simply the last value used by the fitting method.

        Args:
            name(str): Hyper-parameter name.

        Returns:
            float: current value of the specified hyper-parameter.
        """
        flatHyperParameterValues = self._unpackAllHyperParameters(values=True)
        value = flatHyperParameterValues[self._getHyperParameterIndex(self.transitionModel, name)]
        return value

    def getParameterMeanValues(self, name):
        """
        Returns posterior mean values for a parameter of the observation model.

        Args:
            name(str): Name of the parameter to display

        Returns:
            ndarray: array of posterior mean values for the selected parameter
        """
        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        return self.posteriorMeanValues[paramIndex]

    def getParameterDistribution(self, t, name, plot=False, density=True, **kwargs):
        """
        Compute the marginal parameter distribution at a given time step.

        Args:
            t: Time step/stamp for which the parameter distribution is evaluated (or 'avg' for time-averaged parameter
                distribution)
            name(str): Name of the parameter to display
            plot(bool): If True, a plot of the distribution is created
            density(bool): If true, probability density is returned; if false, probability values
            **kwargs: All further keyword-arguments are passed to the plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the parameter values, the second one the corresponding
                probability (density) values
        """
        if self.posteriorSequence == []:
            raise PostProcessingError('Cannot plot posterior sequence as it has not yet been computed. '
                                      'Run complete fit.')

        if t == 'avg':
            # compute time-averaged posterior distribution
            dist = np.sum(self.posteriorSequence, axis=0)/len(self.posteriorSequence)
        else:
            # check if supplied time stamp exists
            if t not in self.formattedTimestamps:
                raise PostProcessingError('Supplied time ({}) does not exist in data or is out of range.'.format(t))
            timeIndex = list(self.formattedTimestamps).index(t)  # to select corresponding posterior distribution

            # select posterior distribution of speciifed time step
            dist = self.posteriorSequence[timeIndex]

        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        axesToMarginalize = list(range(len(self.observationModel.parameterNames)))
        try:
            axesToMarginalize.remove(paramIndex)
        except ValueError:
            raise PostProcessingError('Wrong parameter index. Available indices: {}'.format(axesToMarginalize))

        x = self.marginalGrid[paramIndex]
        dx = self.latticeConstant[paramIndex]
        marginalDistribution = np.squeeze(np.apply_over_axes(np.sum, dist, axesToMarginalize)).copy()

        if density:
            marginalDistribution /= dx

        if plot:
            plt.fill_between(x, 0, marginalDistribution, **kwargs)
            plt.xlabel(self.observationModel.parameterNames[paramIndex])
            if density:
                plt.ylabel('probability density')
            else:
                plt.ylabel('probability')

        return x, marginalDistribution

    def getPD(self, t, name, plot=False, density=True, **kwargs):
        """
        See :meth:`.Study.getParameterDistribution`.
        """
        return self.getParameterDistribution(t, name, plot=plot, density=density, **kwargs)

    def getParameterDistributions(self, name, plot=False, density=True, **kwargs):
        """
        Computes the time series of marginal posterior distributions with respect to a given model parameter.

        Args:
            name(str): Name of the parameter to display
            plot(bool): If True, a plot of the series of distributions is created (density map)
            density(bool): If true, probability density is returned; if false, probability values
            **kwargs: All further keyword-arguments are passed to the plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the parameter values, the second one the sequence of
            corresponding posterior distributions.
        """
        if self.posteriorSequence == []:
            raise PostProcessingError('Cannot plot posterior sequence as it has not yet been computed. '
                                      'Run complete fit.')

        dt = self.formattedTimestamps[1:] - self.formattedTimestamps[:-1]
        if not np.all(dt == dt[0]):
            print('! WARNING: Time stamps are not equally spaced. This may result in false plotting of parameter '
                  'distributions.')

        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        axesToMarginalize = list(range(1, len(self.observationModel.parameterNames) + 1))  # axis 0 is time
        try:
            axesToMarginalize.remove(paramIndex + 1)
        except ValueError:
            raise PostProcessingError('Wrong parameter index. Available indices: {}'
                                      .format(np.array(axesToMarginalize) - 1))

        x = self.marginalGrid[paramIndex]
        dx = self.latticeConstant[paramIndex]
        marginalPosteriorSequence = np.squeeze(
            np.apply_over_axes(np.sum, self.posteriorSequence, axesToMarginalize)).copy()

        if density:
            marginalPosteriorSequence /= dx

        if plot:
            if 'c' in kwargs:
                cmap = createColormap(kwargs['c'])
            elif 'color' in kwargs:
                cmap = createColormap(kwargs['color'])
            else:
                cmap = createColormap('b')

            plt.imshow(marginalPosteriorSequence.T,
                       origin='lower',
                       cmap=cmap,
                       extent=[self.formattedTimestamps[0], self.formattedTimestamps[-1]] + self.boundaries[paramIndex],
                       aspect='auto')

        return x, marginalPosteriorSequence

    def getPDs(self, name, plot=False, density=True, **kwargs):
        """
        See :meth:`.Study.getParameterDistributions`.
        """
        return self.getParameterDistributions(name, plot=plot, density=density, **kwargs)

    def plotParameterEvolution(self, name, color='b', gamma=0.5, **kwargs):
        """
        Extended plot method to display a series of marginal posterior distributions corresponding to a single model
        parameter. In contrast to getMarginalParameterDistributions(), this method includes the removal of plotting
        artefacts, gamma correction as well as an overlay of the posterior mean values.

        Args:
            name(str): name of the parameter to display
            color: color from which a light colormap is created
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5
            kwargs: all further keyword-arguments are passed to the plot of the posterior mean values
        """
        if self.posteriorSequence == []:
            raise PostProcessingError('Cannot plot posterior sequence as it has not yet been computed. '
                                      'Run complete fit.')

        dt = self.formattedTimestamps[1:] - self.formattedTimestamps[:-1]
        if not np.all(dt == dt[0]):
            print('! WARNING: Time stamps are not equally spaced. This may result in false plotting of parameter '
                  'distributions.')

        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        axesToMarginalize = list(range(1, len(self.observationModel.parameterNames) + 1))  # axis 0 is time
        try:
            axesToMarginalize.remove(paramIndex + 1)
        except ValueError:
            raise PostProcessingError('Wrong parameter index to plot. Available indices: {}'
                                      .format(np.array(axesToMarginalize)-1))
        marginalPosteriorSequence = np.squeeze(np.apply_over_axes(np.sum, self.posteriorSequence, axesToMarginalize))

        # clean up very small probability values, as they may create image artefacts
        pmax = np.amax(marginalPosteriorSequence)
        marginalPosteriorSequence[marginalPosteriorSequence < pmax*(10**-20)] = 0

        plt.imshow(marginalPosteriorSequence.T**gamma,
                   origin='lower',
                   cmap=createColormap(color),
                   extent=[self.formattedTimestamps[0], self.formattedTimestamps[-1]] + self.boundaries[paramIndex],
                   aspect='auto')

        # set default color of plot to black
        if ('c' not in kwargs) and ('color' not in kwargs):
            kwargs['c'] = 'k'

        # set default linewidth to 1.5
        if ('lw' not in kwargs) and ('linewidth' not in kwargs):
            kwargs['lw'] = 1.5

        plt.plot(self.formattedTimestamps, self.posteriorMeanValues[paramIndex], **kwargs)

        plt.ylim(self.boundaries[paramIndex])
        plt.ylabel(self.observationModel.parameterNames[paramIndex])
        plt.xlabel('time step')

    def plot(self, name, **kwargs):
        """
        Convenience method to plot the temporal evolution of observation model parameters, or the parameter distribution
        at a specific time step. Extended functionality for other study classes.

        Args:
            name(str): name of the parameter to display
            color: color from which a light colormap is created (for parameter evolution only)
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5 (for
                parameter evolution only)
            t: Time step/stamp for which the parameter distribution is evaluated
            density(bool): If true, probability density is plotted; if false, probability values
            kwargs: all further keyword-arguments are passed to the axes object of the plot
        """
        density = kwargs.pop('density', True)

        # plot parameter distribution at specific time step
        if 't' in kwargs.keys():
            t = kwargs.pop('t')
            self.getParameterDistribution(t, name, plot=True, density=density, **kwargs)
        # plot parameter evolution
        else:
            # read additional kwargs (or set default values)
            color = kwargs.pop('color', 'b')
            gamma = kwargs.pop('gamma', 0.5)
            self.plotParameterEvolution(name, color=color, gamma=gamma, **kwargs)

    def _checkConsistency(self):
        """
        This method is called at the very beginning of analysis methods to ensure that all necessary elements of the
        model are set correctly. If problem with user input is detected, an exception will be raised.
        """
        if len(self.rawData) == 0:
            raise ConfigurationError('No data loaded.')
        if not self.observationModel:
            raise ConfigurationError('No observation model chosen.')
        if not self.transitionModel:
            raise ConfigurationError('No transition model chosen.')

        # check for duplicate hyper-parameter names
        flatNames = self._unpackAllHyperParameters(values=False)
        u, i = np.unique(flatNames, return_inverse=True)
        duplicates = u[np.bincount(i) > 1]
        if len(duplicates) > 0:
            raise ConfigurationError('Detected duplicate hyper-parameter names: {}.'.format(duplicates))


class HyperStudy(Study):
    """
    Infers hyper-parameter distributions. This class serves as an extension to the basic Study class and allows to
    compute the distribution of hyper-parameters of a given transition model. For further information, see the
    documentation of the fit-method of this class.
    """
    def __init__(self, silent=False):
        super(HyperStudy, self).__init__(silent=silent)

        self.hyperGrid = []
        self.hyperGridValues = []
        self.hyperGridConstant = []
        self.flatHyperParameters = []
        self.flatHyperParameterNames = []
        self.flatHyperPriors = []
        self.flatHyperPriorValues = []
        self.hyperParameterDistribution = None
        self.averagePosteriorSequence = None
        self.logEvidenceList = []
        self.localEvidenceList = []

        if not silent:
            print('  --> Hyper-study')

    def _createHyperGrid(self, silent=False):
        """
        Creates an array of hyper-parameter values that are fitted. Also determines grid constants for proper
        normalisation and computes prior probability (density) values

        Args:
            silent(bool): If true, no output is produced by this method
        """
        # extract flat list of hyper-parameter names and values
        self.flatHyperParameters = self._unpackAllHyperParameters()
        self.flatHyperParameterNames = self._unpackAllHyperParameters(values=False)
        self.flatHyperPriors = self._unpackAllHyperPriors()

        # look if change/break-points have value 'all' and assign array of all time-stamps
        for i, v in enumerate(self.flatHyperParameters):
            if isinstance(v, str) and v == 'all':
                self.flatHyperParameters[i] = self.formattedTimestamps[:-1]

        # create hyper-parameter grid
        temp = np.meshgrid(*self.flatHyperParameters, indexing='ij')
        if len(self.flatHyperParameterNames) > 0:
            self.hyperGridValues = np.array([t.ravel() for t in temp]).T
        else:
            self.hyperGridValues = np.array([])

        # find lattice constants for equally spaced hyper-parameter values
        self.hyperGridConstant = []
        for values in self.flatHyperParameters:
            if isinstance(values, Iterable) and len(values) > 1:
                a = np.array(values)
                d = a[1:] - a[:-1]
                dd = d[1:] - d[:-1]
                if np.all(np.abs(dd) < 10 ** -10):  # for equally spaced values, set difference as grid-constant
                    self.hyperGridConstant.append(np.abs(d[0]))
                else:  # for irregularly spaced values (e.g. categorical), set grid-constant to 1
                    self.hyperGridConstant.append(1)
            else:  # for single value, set grid-constant to 1
                self.hyperGridConstant.append(1)
        self.hyperGridConstant = np.array(self.hyperGridConstant)

        # evaluate hyper-prior values
        priorValuesList = []
        priorNamesList = []
        for prior, values, gridConst, name in zip(self.flatHyperPriors, self.flatHyperParameters,
                                                  self.hyperGridConstant, self.flatHyperParameterNames):
            if prior is None:  # no prior specified
                priorValues = np.ones_like(values, dtype=float)
                priorValues /= np.sum(priorValues)
                priorValues /= gridConst
                priorNamesList.append('uniform')

            elif hasattr(prior, '__call__'):  # prior specified by function
                try:
                    priorValues = np.array([prior(value) for value in values])
                    norm = np.sum(priorValues)
                    priorValues /= norm
                    priorValues /= gridConst
                except:
                    raise ConfigurationError('Failed to set hyper-prior for "{}" from function "{}".'
                                             .format(name, prior.__name__))
                if norm != 1.:
                    priorNamesList.append(prior.__name__ + ' (re-normalized)')
                else:
                    priorNamesList.append(prior.__name__)

            elif isinstance(prior, Iterable):  # prior specified by list/array
                if len(prior) != len(values):
                    raise ConfigurationError('Failed to set hyper-prior for "{}" from list/array.'.format(name))
                priorValues = prior
                norm = np.sum(priorValues)
                priorValues /= norm
                priorValues /= gridConst
                if norm != 1.:
                    priorNamesList.append('list/array (re-normalized)')
                else:
                    priorNamesList.append('list/array')

            else:  # SymPy RV
                if len(freeSymbols(prior)) > 0:
                    raise ConfigurationError('Hyper-prior for "{}" must not contain free parameters.'.format(name))

                # get symbolic representation of probability density
                x = abc.x
                symDensity = density(prior)(x)

                # get density as lambda function
                pdf = lambdify([x], symDensity, modules=['numpy', {'factorial': factorial}])

                # evaluate density
                priorValues = pdf(values)
                priorNamesList.append(str(symDensity))

            priorValuesList.append(priorValues)

        # create hyper-prior grid
        if len(self.flatHyperParameterNames) > 0:
            temp = np.meshgrid(*priorValuesList, indexing='ij')
            self.flatHyperPriorValues = np.array([t.ravel() for t in temp]).T
            self.flatHyperPriorValues = np.prod(self.flatHyperPriorValues, axis=1)  # joint prob. density
            if not silent and len(self.hyperGridValues) > 1:
                print('+ Set hyper-prior(s): {}'.format(priorNamesList))
        else:
            # we need a dummy value for transition models without hyper-parameters
            self.flatHyperPriorValues = np.array([1])

    def fit(self, forwardOnly=False, evidenceOnly=False, silent=False, nJobs=1, customHyperGrid=False):
        """
        This method over-rides the according method of the Study-class. It runs the algorithm for equally spaced hyper-
        parameter values as defined by the variable 'hyperGrid'. The posterior sequence represents the average
        model of all analyses. Posterior mean values are computed from this average model.

        Args:
            forwardOnly(bool): If set to True, the fitting process is terminated after the forward pass. The resulting
                posterior distributions are so-called "filtering distributions" which - at each time step -
                only incorporate the information of past data points. This option thus emulates an online
                analysis.
            evidenceOnly(bool): If set to True, only forward pass is run and evidence is calculated. In contrast to the
                forwardOnly option, no posterior mean values are computed and no posterior distributions are stored.
            silent(bool): If set to true, reduced output is created by this method.
            nJobs(int): Number of processes to employ. Multiprocessing is based on the 'pathos' module.
            customHyperGrid(bool): If set to true, the method "_createHyperGrid" is not called before starting the fit.
                This is used by the class "ChangepointStudy", which employs a custom version of "_createHyperGrid".
        """
        self.fitWarningCounter = 0

        # format data/timestamps once, so number of data segments is known and _createGrid() works properly
        self.formattedData = movingWindow(self.rawData, self.observationModel.segmentLength)
        self.formattedTimestamps = self.rawTimestamps[self.observationModel.segmentLength - 1:]

        # create hyper-parameter grid
        if not customHyperGrid:
            self._createHyperGrid(silent=silent)

            # additional consistency check if multiple change/break-points share identical values.
            # in this case, a ChangepointStudy should be used!
            names = list(flatten(self._unpackChangepointNames(self.transitionModel))) + \
                    list(flatten(self._unpackBreakpointNames(self.transitionModel)))
            if len(names) > 1:
                indices = [self.flatHyperParameterNames.index(name) for name in names]
                values = self.hyperGridValues[:, indices]

                for v in values:
                    if np.unique(v).size < v.size:
                        raise ConfigurationError('Detected multiple change-/break-points with identical values and/or '
                                                 'overlapping value intervals. Use "ChangepointStudy" instead of '
                                                 '"HyperStudy" for such cases.')

        self._checkConsistency()

        if not evidenceOnly:
            # The average posterior distribution is stored in log-space for numerical stability. This way, it can be
            # updated iteratively without storing all individual posterior distributions (after the iteration over all
            # hyper-parameter values is done, it is transformed back to linear space)
            self.averagePosteriorSequence = np.zeros([len(self.formattedData)]+self.gridSize) - np.inf

        self.logEvidenceList = []
        self.localEvidenceList = []

        # hyper-study fit is only necessary for more than one combination of hyper-parameter values
        if len(self.hyperGridValues) > 1:
            if not silent:
                print('+ Started new fit.')
                print('    + {} analyses to run.'.format(len(self.hyperGridValues)))

            # check if multiprocessing is available
            if nJobs > 1:
                try:
                    from pathos.multiprocessing import ProcessPool
                except ImportError:
                    raise ImportError('No module named pathos.multiprocessing. This module represents an optional '
                                      'dependency of bayesloop and is therefore not installed alongside bayesloop.')

                # prepare parallel execution
                if not silent:
                    print('    + Creating {} processes.'.format(nJobs))
                pool = ProcessPool(nodes=nJobs)

                # use parallelFit method to create copies of this HyperStudy instance with only partial
                # hyper-grid values
                subStudies = pool.map(self._parallelFit,
                                      range(nJobs),
                                      [nJobs]*nJobs,
                                      [forwardOnly]*nJobs,
                                      [evidenceOnly]*nJobs,
                                      [silent]*nJobs)

                # prevent memory pile-up in main process
                pool.close()
                pool.join()
                pool.terminate()
                pool.restart()

                # merge all sub-studies
                for S in subStudies:
                    self.logEvidenceList += S.logEvidenceList
                    self.localEvidenceList += S.localEvidenceList
                    if not evidenceOnly:
                        self.averagePosteriorSequence = np.logaddexp(self.averagePosteriorSequence,
                                                                     S.averagePosteriorSequence)
            # single process fit
            else:
                # show progressbar if silent=False
                if not silent:
                    enum = tqdm(enumerate(self.hyperGridValues), total=len(self.hyperGridValues))
                else:
                    enum = enumerate(self.hyperGridValues)

                for i, hyperParamValues in enum:
                    self._setSelectedHyperParameters(hyperParamValues)

                    # call fit method from parent class
                    Study.fit(self, forwardOnly=forwardOnly, evidenceOnly=evidenceOnly, silent=True)

                    self.logEvidenceList.append(self.logEvidence)
                    self.localEvidenceList.append(self.localEvidence)

                    if (not evidenceOnly) and np.isfinite(self.logEvidence):
                        # For numerical stability, zeros in posterior distribution are replaced with small constant.
                        # Note that this is typically only the case for hyper-parameter values with low likelihood. The
                        # procedure therefore has no (measurable) effect on the average posterior sequence.
                        self.posteriorSequence[self.posteriorSequence < 10.**-300] = 10.**-300
                        self.averagePosteriorSequence = np.logaddexp(self.averagePosteriorSequence,
                                                                     np.log(self.posteriorSequence) +
                                                                     self.logEvidence +
                                                                     np.log(self.flatHyperPriorValues[i]))

                # remove progressbar correctly
                if not silent:
                    enum.close()

            if not evidenceOnly:
                # transform average posterior distribution into linear space
                # (we prefer underflows rather than overflows)
                self.averagePosteriorSequence -= np.amax(self.averagePosteriorSequence)
                self.averagePosteriorSequence = np.exp(self.averagePosteriorSequence)

                # compute average posterior distribution
                normalization = np.array([np.sum(posterior) for posterior in self.averagePosteriorSequence])
                for i in range(len(self.grid)):
                    normalization = normalization[:, None]  # add axis; needs to match averagePosteriorSequence
                self.averagePosteriorSequence /= normalization

                # set self.posteriorSequence to average posterior sequence for plotting reasons
                self.posteriorSequence = self.averagePosteriorSequence

                if not silent:
                    print('    + Computed average posterior sequence')

            # compute hyper-parameter distribution
            logHyperParameterDistribution = self.logEvidenceList + np.log(self.flatHyperPriorValues) + \
                                            np.sum(np.log(self.hyperGridConstant))

            # transform hyper-parameter distribution into linear space
            # (we prefer underflows rather than overflows)
            scaledLogHyperParameterDistribution = logHyperParameterDistribution - np.amax(logHyperParameterDistribution)
            self.hyperParameterDistribution = np.exp(scaledLogHyperParameterDistribution)
            self.hyperParameterDistribution /= np.sum(self.hyperParameterDistribution)
            self.hyperParameterDistribution /= np.prod(self.hyperGridConstant)  # probability density

            if not silent:
                print('    + Computed hyper-parameter distribution')

            # compute log-evidence of average model
            self.logEvidence = logsumexp(logHyperParameterDistribution)
            if not silent:
                print('    + Log10-evidence of average model: {:.5f}'.format(self.logEvidence / np.log(10)))

            # compute local evidence of average model
            self.localEvidence = np.sum((np.array(self.localEvidenceList).T*self.flatHyperPriorValues).T, axis=0)

            if not silent:
                print('    + Computed local evidence of average model')

            # compute posterior mean values
            if not evidenceOnly:
                self.posteriorMeanValues = np.empty([len(self.grid), len(self.posteriorSequence)])
                for i in range(len(self.grid)):
                    self.posteriorMeanValues[i] = np.array([np.sum(p*self.grid[i]) for p in self.posteriorSequence])

                if not silent:
                    print('    + Computed mean parameter values.')

            # clear localEvidenceList (to keep file size small for stored studies)
            self.localEvidenceList = []

            # restore hyper-parameter values (individual values have been set during fitting)
            self._setAllHyperParameters(self.flatHyperParameters)

            if not silent:
                print('+ Finished fit.')

        # only one combination of hyper-parameter values
        else:
            if not silent:
                if len(self.hyperGridValues) == 1:
                    print('+ Only one combination of hyper-parameter values, switching to standard fit method.')
                if len(self.hyperGridValues) == 0:
                    print('+ Transition model contains no hyper-parameters, switching to standard fit method.')

            Study.fit(self, forwardOnly=forwardOnly, evidenceOnly=evidenceOnly, silent=silent)

    def _parallelFit(self, idx, nJobs, forwardOnly, evidenceOnly, silent):
        """
        This method is called by the fit method of the HyperStudy class. It creates a copy of the current class
        instance and performs a fit based on a subset of the specified hyper-parameter grid. The method thus allows
        to distribute a HyperStudy fit among multiple processes for multiprocessing.

        Args:
            idx(int): Index from 0 to (nJobs-1), indicating which part of the hyper-grid values are to be analyzed.
            nJobs(int): Number of processes to employ. Multiprocessing is based on the 'pathos' module.
            forwardOnly(bool): If set to True, the fitting process is terminated after the forward pass. The resulting
                posterior distributions are so-called "filtering distributions" which - at each time step -
                only incorporate the information of past data points. This option thus emulates an online
                analysis.
            evidenceOnly(bool): If set to True, only forward pass is run and evidence is calculated. In contrast to the
                forwardOnly option, no posterior mean values are computed and no posterior distributions are stored.
            silent(bool): If set to True, no output is generated by the fitting method.

        Returns:
            HyperStudy instance
        """
        S = copy(self)
        S.hyperGridValues = np.array_split(S.hyperGridValues, nJobs)[idx]
        S.flatHyperPriorValues = np.array_split(S.flatHyperPriorValues, nJobs)[idx]

        # show progressbar for last process if silent=False
        if not silent and idx == nJobs-1:
            enum = tqdm(enumerate(S.hyperGridValues), total=len(S.hyperGridValues))
        else:
            enum = enumerate(S.hyperGridValues)

        for i, hyperParamValues in enum:
            S._setSelectedHyperParameters(hyperParamValues)

            # call fit method from parent class
            Study.fit(S, forwardOnly=forwardOnly, evidenceOnly=evidenceOnly, silent=True)

            S.logEvidenceList.append(S.logEvidence)
            S.localEvidenceList.append(S.localEvidence)
            if (not evidenceOnly) and np.isfinite(S.logEvidence):
                # For numerical stability, zeros in posterior distribution are replaced with small constant.
                # Note that this is typically only the case for hyper-parameter values with low likelihood. The
                # procedure therefore has no (measurable) effect on the average posterior sequence.
                S.posteriorSequence[S.posteriorSequence < 10. ** -300] = 10. ** -300
                S.averagePosteriorSequence = np.logaddexp(S.averagePosteriorSequence,
                                                          np.log(S.posteriorSequence) +
                                                          S.logEvidence +
                                                          np.log(S.flatHyperPriorValues[i]))

        # remove progressbar correctly
        if not silent and idx == nJobs-1:
            enum.close()

        return S

    # optimization methods are inherited from Study class, but cannot be used in this case
    def optimize(self, *args, **kwargs):
        raise NotImplementedError('HyperStudy object has no optimizing method.')

    def _unpackHyperPriors(self, transitionModel):
        """
        Returns list of all hyper-priors, nested as the transition model.

        Args:
            transitionModel: An instance of a transition model

        Returns:
            list: hyper-priors
        """
        priorList = []
        # recursion step for sub-models
        if hasattr(transitionModel, 'models'):
            for m in transitionModel.models:
                priorList.append(self._unpackHyperPriors(m))

        # append prior
        if hasattr(transitionModel, 'prior'):
            # only take prior if transition model has at least one hyper-parameter or is a break-point
            # otherwise, the number of hyper-priors does not match the number of hyper-parameters
            if (hasattr(transitionModel, 'hyperParameterNames') and len(transitionModel.hyperParameterNames) > 0) or\
                    (str(transitionModel) == 'Break-point'):
                priorList.append(transitionModel.prior)
        return priorList

    def _unpackAllHyperPriors(self):
        """
        Returns a flattened list of all hyper-priors of the current transition model.

        Returns:
            list: all hyper-priors of the current transition model
        """
        return list(flatten(self._unpackHyperPriors(self.transitionModel)))

    def getHyperParameterDistribution(self, name, plot=False, **kwargs):
        """
        Computes marginal hyper-parameter distribution of a single hyper-parameter in a HyperStudy fit.

        Args:
            name(str): Name of the hyper-parameter to display
                (first model hyper-parameter)
            plot(bool): If True, a bar chart of the distribution is created
            **kwargs: All further keyword-arguments are passed to the bar-plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the hyper-parameter values, the second one the
                corresponding probability values
        """
        # check if only a standard fit has been carried out
        if len(self.hyperGridValues) < 2:
            raise PostProcessingError('At least two combinations of hyper-parameter values need to be fitted to '
                                      'evaluate a hyper-parameter distribution. Check transition model.')

        paramIndex = self._getHyperParameterIndex(self.transitionModel, name)

        axesToMarginalize = list(range(len(self.flatHyperParameterNames)))
        axesToMarginalize.remove(paramIndex)

        # reshape hyper-parameter distribution for easy marginalizing
        hyperGridSteps = []
        for x in self.flatHyperParameters:
            if isinstance(x, Iterable):
                hyperGridSteps.append(len(x))
            else:
                hyperGridSteps.append(1)

        distribution = self.hyperParameterDistribution.reshape(hyperGridSteps, order='C')
        marginalDistribution = np.squeeze(np.apply_over_axes(np.sum, distribution, axesToMarginalize))
        marginalDistribution *= np.prod(self.hyperGridConstant)  # convert to probability (from density)

        x = self.flatHyperParameters[paramIndex]
        if plot:
            # check if categorical
            if np.any(np.abs(np.diff(np.diff(x))) > 10 ** -10):
                plt.bar(np.arange(len(x)), marginalDistribution, align='center', width=1., **kwargs)
                plt.xticks(np.arange(len(x)), x)
                plt.ylabel('probability')
            # regular spacing
            else:
                plt.bar(x, marginalDistribution, align='center', width=self.hyperGridConstant[paramIndex], **kwargs)
                plt.ylabel('probability')

            plt.xlabel(self.flatHyperParameterNames[paramIndex])

        return x, marginalDistribution

    def getHPD(self, name, plot=False, **kwargs):
        """
        See :meth:`.HyperStudy.getHyperParameterDistribution`.
        """
        return self.getHyperParameterDistribution(name, plot=plot, **kwargs)

    def getJointHyperParameterDistribution(self, names, plot=False, figure=None, subplot=111, **kwargs):
        """
        Computes the joint distribution of two hyper-parameters of a HyperStudy and optionally creates a 3D bar chart.
        Note that the 3D plot can only be included in an existing plot by passing a figure object and subplot
        specification.

        Args:
            names(list): List of two hyper-parameter names to display
            plot(bool): If True, a 3D-bar chart of the distribution is created
            figure: In case the plot is supposed to be part of an existing figure, it can be passed to the method. By
                default, a new figure is created.
            subplot: Characterization of subplot alignment, as in matplotlib. Default: 111
            **kwargs: all further keyword-arguments are passed to the bar3d-plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray, ndarray: The first and second array contains the hyper-parameter values, the
                third one the corresponding probability values
        """
        # check if only a standard fit has been carried out
        if len(self.hyperGridValues) < 2:
            raise PostProcessingError('At least two combinations of hyper-parameter values need to be fitted to '
                                      'evaluate a hyper-parameter distribution. Check transition model.')

        # check if list with two elements is provided
        if not isinstance(names, Iterable):
            raise PostProcessingError('A list of exactly two hyper-parameters has to be provided.')
        elif not len(names) == 2:
            raise PostProcessingError('A list of exactly two hyper-parameters has to be provided.')

        paramIndices = [self._getHyperParameterIndex(self.transitionModel, n) for n in names]

        # put parameter indices in ascending order for marginalization and plotting
        # original order is restored for returned values
        switch = False
        if not paramIndices[0] < paramIndices[1]:
            switch = True
            paramIndices = paramIndices[::-1]

        axesToMarginalize = list(range(len(self.flatHyperParameterNames)))
        for p in paramIndices:
            axesToMarginalize.remove(p)

        # reshape hyper-parameter distribution for easy marginalizing
        hyperGridSteps = []
        for x in self.flatHyperParameters:
            if isinstance(x, Iterable):
                hyperGridSteps.append(len(x))
            else:
                hyperGridSteps.append(1)

        distribution = self.hyperParameterDistribution.reshape(hyperGridSteps, order='C')
        marginalDistribution = np.squeeze(np.apply_over_axes(np.sum, distribution, axesToMarginalize))
        marginalDistribution *= np.prod(self.hyperGridConstant)  # convert to probability (from density)

        x, y = [self.flatHyperParameters[i] for i in paramIndices]
        if np.any(np.abs(np.diff(np.diff(x))) > 10 ** -10):
            x2 = np.tile(np.arange(len(x)), (len(y), 1)).T
        else:
            x2 = np.tile(x, (len(y), 1)).T

        if np.any(np.abs(np.diff(np.diff(y))) > 10 ** -10):
            y2 = np.tile(np.arange(len(y)), (len(x), 1))
        else:
            y2 = np.tile(y, (len(x), 1))

        z = marginalDistribution

        if plot:
            # allow to add plot to predefined figure
            if figure is None:
                fig = plt.figure()
            else:
                fig = figure
            ax = fig.add_subplot(subplot, projection='3d')

            ax.bar3d(x2.flatten() - self.hyperGridConstant[paramIndices[0]]/2.,
                     y2.flatten() - self.hyperGridConstant[paramIndices[1]]/2.,
                     z.flatten()*0.,
                     self.hyperGridConstant[paramIndices[0]],
                     self.hyperGridConstant[paramIndices[1]],
                     z.flatten(),
                     zsort='max',
                     **kwargs
                     )

            # check for categorical hyper-parameter values
            if np.any(np.abs(np.diff(np.diff(x))) > 10 ** -10):
                plt.xticks(np.arange(len(x)), x)
            if np.any(np.abs(np.diff(np.diff(y))) > 10 ** -10):
                plt.yticks(np.arange(len(y)), y)

            ax.set_xlabel(self.flatHyperParameterNames[paramIndices[0]])
            ax.set_ylabel(self.flatHyperParameterNames[paramIndices[1]])

            ax.set_zlabel('probability')

        # restore original order of parameters before returning values
        if switch:
            x, y = y, x
            marginalDistribution = marginalDistribution.T

        return x, y, marginalDistribution

    def getJHPD(self, names, plot=False, figure=None, subplot=111, **kwargs):
        """
        See :meth:`.HyperStudy.getJointHyperParameterDistribution`.
        """
        return self.getJointHyperParameterDistribution(names, plot=plot, figure=figure, subplot=subplot, **kwargs)

    def plot(self, name, **kwargs):
        """
        Convenience method to plot the temporal evolution of observation model parameters, the distribution of a
        parameter at a specific time step, or the distribution of a hyper-parameter.

        Args:
            name(str): name of the (hyper-)parameter to display
            color: color from which a light colormap is created (for parameter evolution only)
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5 (for
                parameter evolution only)
            t: Time step/stamp for which the parameter distribution is evaluated
            density(bool): If true, probability density is plotted; if false, probability values. Note: Only availble
                for parameters, not hyper-parameters.
            kwargs: all further keyword-arguments are passed to the axes object of the plot
        """
        density = kwargs.pop('density', True)

        # plot parameter distribution at specific time step
        if 't' in kwargs.keys():
            t = kwargs.pop('t')
            self.getParameterDistribution(t, name, plot=True, density=density, **kwargs)

        else:
            # check is name belongs to hyper-parameter
            try:
                self._getHyperParameterIndex(self.transitionModel, name)
                hyper = True
            except PostProcessingError:
                hyper = False

            # for hyper-parameters, plot distribution
            if hyper:
                self.getHyperParameterDistribution(name, plot=True, **kwargs)
            # for parameters, plot temporal evolution
            else:
                # read additional kwargs (or set default values)
                color = kwargs.pop('color', 'b')
                gamma = kwargs.pop('gamma', 0.5)
                self.plotParameterEvolution(name, color=color, gamma=gamma, **kwargs)


class ChangepointStudy(HyperStudy):
    """
    Infers change-points and structural breaks. This class builds on the HyperStudy-class and the change-point
    transition model to perform a series of analyses with varying change point times. It subsequently computes the
    average model from all possible change points and creates a probability distribution of change point times. It
    supports any number of change-points and arbitarily combined models.
    """
    def __init__(self, silent=False):
        super(ChangepointStudy, self).__init__(silent=silent)

        # store all possible combinations of change-points (even the ones that are assigned a probability of zero),
        # to reconstruct change-point distribution after analysis
        self.allHyperGridValues = []
        self.allHyperPriorValues = []
        self.mask = []  # mask to select valid change-point combinations

        self.userDefinedGrid = False  # needed to ensure that user-defined hyper-grid is not overwritten by fit-method
        self.hyperGridBackup = []  # needed to reconstruct hyperGrid attribute in the case of break-point model

        if not silent:
            print('  --> Change-point analysis')

    def fit(self, forwardOnly=False, evidenceOnly=False, silent=False, nJobs=1):
        """
        This method over-rides the corresponding method of the HyperStudy-class. It runs the algorithm for all possible
        combinations of change-points (and possible scans a range of values for other hyper-parameters). The posterior
        sequence represents the average model of all analyses. Posterior mean values are computed from this average
        model.

        Args:
            forwardOnly(bool): If set to True, the fitting process is terminated after the forward pass. The resulting
                posterior distributions are so-called "filtering distributions" which - at each time step -
                only incorporate the information of past data points. This option thus emulates an online
                analysis.
            evidenceOnly(bool): If set to True, only forward pass is run and evidence is calculated. In contrast to the
                forwardOnly option, no posterior mean values are computed and no posterior distributions are stored.
            silent(bool): If set to True, reduced output is generated by the fitting method.
            nJobs(int): Number of processes to employ. Multiprocessing is based on the 'pathos' module.
        """
        # format data/timestamps once, so number of data segments is known
        self.formattedData = movingWindow(self.rawData, self.observationModel.segmentLength)
        self.formattedTimestamps = self.rawTimestamps[self.observationModel.segmentLength - 1:]

        # nested serial transition models are not supported, as the correct order is not determined correctly
        if len(list(flatten(self._unpackSerialTransitionModels(self.transitionModel)))) > 1:
            raise NotImplementedError('Multiple instances of SerialTransition models are currently not supported by '
                                      'ChangepointStudy.')

        # determine names of change/break-points
        changepoints = list(flatten(self._unpackChangepointNames(self.transitionModel)))
        breakpoints = list(flatten(self._unpackBreakpointNames(self.transitionModel)))

        # both types are not allowed at the moment, as the correct order is not determined correctly
        if len(changepoints) > 0 and len(breakpoints) > 0:
            raise NotImplementedError('Detected both change-points (Changepoint transition model) and break-points '
                                      '(SerialTransitionModel). Currently, only one type is supported in a single '
                                      'transition model.')

        # at least one change/break-point should be present
        if len(changepoints) == 0 and len(breakpoints) == 0:
            raise ConfigurationError('No change-points or break-points detected in transition model. Check transition '
                                     'model.')

        self.flatHyperParameters = self._unpackAllHyperParameters()
        self.flatHyperParameterNames = self._unpackAllHyperParameters(values=False)

        # create hyperGrid in the case of change-points
        if len(changepoints) > 0:
            print('+ Detected {} change-point(s) in transition model: {}'.format(len(changepoints), changepoints))
            points = changepoints
        else:
            print('+ Detected {} break-point(s) in transition model: {}'.format(len(breakpoints), breakpoints))
            points = breakpoints

        # first create standard hyper-grid
        self._createHyperGrid()
        self.allHyperGridValues = self.hyperGridValues[:]
        self.allHyperPriorValues = self.flatHyperPriorValues[:]

        # extract hyper-grid values that belong to changepoints
        pointMask = np.sum([np.array(self.flatHyperParameterNames) == p for p in points], axis=0).astype(bool)
        maskedHyperGridValues = self.allHyperGridValues[:, pointMask]

        # only accept if change-point values are ordered (and not equal)
        self.mask = np.array(
            [all(x[i] < x[i + 1] for i in range(len(points) - 1)) for x in maskedHyperGridValues],
            dtype=bool)
        self.hyperGridValues = self.allHyperGridValues[self.mask]
        self.flatHyperPriorValues = self.allHyperPriorValues[self.mask]

        # correct prior values to ensure correct normalization after sorting out combinations
        self.flatHyperPriorValues *= np.sum(self.allHyperPriorValues)/np.sum(self.allHyperPriorValues[self.mask])

        # call fit method of hyper-study
        HyperStudy.fit(self,
                       forwardOnly=forwardOnly,
                       evidenceOnly=evidenceOnly,
                       silent=silent,
                       nJobs=nJobs,
                       customHyperGrid=True)

        # for proper plotting, hyperGridValues must include all possible combinations of hyper-parameter values. We
        # therefore have to include invalid combinations and assign the probability zero to them.
        temp = np.zeros(len(self.allHyperGridValues))
        temp[self.mask] = self.hyperParameterDistribution
        self.hyperParameterDistribution = temp

        temp = np.zeros(len(self.allHyperPriorValues))
        temp[self.mask] = self.flatHyperPriorValues
        self.flatHyperPriorValues = temp

    def _unpackSerialTransitionModels(self, transitionModel):
        """
        Returns list of all occurrences of serial transition models in the transition model, nested like the transition
        model.

        Returns:
            list: all serial transition models
        """
        modelList = []
        # recursion step for sub-models
        if hasattr(transitionModel, 'models'):
            for m in transitionModel.models:
                modelList.append(self._unpackSerialTransitionModels(m))

        # extend hyper-parameter based on current (sub-)model
        if hasattr(transitionModel, 'hyperParameterNames'):
            if str(transitionModel) == 'Serial transition model':
                modelList.append(transitionModel)

        return modelList

    def getDurationDistribution(self, names, plot=False, **kwargs):
        """
        Computes the distribution of the number of time steps between two change/break-points. This distribution of
        duration is created from the joint distribution of the two specified change/break-points.

        Args:
            names(list): List of two parameter names of change/break-points to display
                (first and second model parameter)
            plot(bool): If True, a bar chart of the distribution is created
            **kwargs: All further keyword-arguments are passed to the bar-plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the number of time steps, the second one the corresponding
                probability values.
        """
        # check if list with two elements is provided
        if not isinstance(names, Iterable):
            raise PostProcessingError('A list of exactly two hyper-parameters has to be provided.')
        elif not len(names) == 2:
            raise PostProcessingError('A list of exactly two hyper-parameters has to be provided.')
        paramIndices = [self._getHyperParameterIndex(self.transitionModel, n) for n in names]

        # check if parameter indices are in ascending order (so axes are labeled correctly)
        if not paramIndices[0] < paramIndices[1]:
            print('! WARNING: Switching order of change-/breakpoints to obtain positive values for duration.')
            paramIndices = paramIndices[::-1]

        values = self.hyperGridValues[:, paramIndices].T
        duration = np.unique(values[1] - values[0])  # get all possible differences between time points
        durationDistribution = np.zeros(len(duration))  # initialize array for distribution

        # loop over all hyper-grid points and collect probabilities for different durations
        for i, values in enumerate(self.allHyperGridValues[:, paramIndices]):
            if values[1] > values[0]:
                # get matching index in duration (rounding needed because of finite precision)
                idx = np.where(duration.round(10) == (values[1]-values[0]).round(10))[0][0]
                durationDistribution[idx] += self.hyperParameterDistribution[i]

        # properly normalize duration distribution
        durationDistribution /= np.sum(durationDistribution)

        if plot:
            plt.bar(duration, durationDistribution, align='center', width=duration[0], **kwargs)

            plt.xlabel('duration between {} and {} (in time steps)'
                       .format(self.flatHyperParameterNames[paramIndices[0]],
                               self.flatHyperParameterNames[paramIndices[1]]))
            plt.ylabel('probability')

        return duration, durationDistribution

    def getDD(self, names, plot=False, **kwargs):
        """
        See :meth:`.ChangepointStudy.getDurationDistribution`.
        """
        return self.getDurationDistribution(names, plot=plot, **kwargs)


class OnlineStudy(HyperStudy):
    """
    Enables model selection for online data streams. This class builds on the Study-class and features a step-method
    to include new data points in the study as they arrive from a data stream. This online-analysis is performed in an
    forward-only way, resulting in filtering-distributions only. In contrast to a normal study, however, one can add
    multiple transition models to account for different types of parameter dynamics (similar to a Hyper study). The
    Online study then computes the probability distribution over all transition models for each new data point (or all
    past data points), enabling real-time model selection.

    Args:
        storeHistory(bool): If true, posterior distributions and their mean values, as well as hyper-posterior
            distributions are stored for all time steps.
    """
    def __init__(self, storeHistory=False, silent=False):
        super(OnlineStudy, self).__init__(silent=silent)

        self.firstStep = True

        self.transitionModels = []
        self.transitionModelNames = []
        self.tmCount = None
        self.tmCounts = []
        self.hyperParameterValues = []
        self.allFlatHyperParameterValues = []
        self.hyperParameterNames = []
        self.hyperGridConstants = []

        self.logEvidenceList = None
        self.hyperLogEvidenceList = None
        self.hyperPrior = []
        self.hyperPriorValues = []
        self.transitionModelPrior = None

        self.parameterPosterior = None
        self.transitionModelPosterior = None
        self.marginalizedPosterior = None

        self.hyperParameterDistribution = None
        self.transitionModelDistribution = None
        self.localTransitionModelDistribution = None

        self.storeHistory = storeHistory
        self.posteriorMeanValues = []
        self.posteriorSequence = []
        self.hyperParameterSequence = []
        self.transitionModelSequence = []
        self.localTransitionModelSequence = []

        if not silent:
            print('  --> Online study')

    def addTransitionModel(self, name, transitionModel):
        """
        Adds a transition model to the list of transition models that are fitted in each time step. Note that a list of
        hyper-parameter values can be supplied.

        Args:
            name(str): a custom name for this transition model to identify it in post-processing methods
            transitionModel: instance of a transition model class.

        Example:
            Here, 'S' denotes the OnlineStudy instance. In the first example, we assume a Poisson observation model and
            add a Gaussian random walk with varying standard deviation to the rate parameter 'lambda':

                S.setObservationModel(bl.om.Poisson('lambda', bl.oint(0, 6, 1000)))
                S.addTransitionModel(bl.tm.GaussianRandomWalk('sigma', [0, 0.1, 0.2, 0.3], target='lambda'))
        """
        # extract hyper-parameter values and names
        self.setTransitionModel(transitionModel, silent=True)
        self._createHyperGrid(silent=True)

        self.transitionModels.append(transitionModel)
        self.transitionModelNames.append(name)
        self.hyperParameterValues.append(self.hyperGridValues[:])
        self.allFlatHyperParameterValues.append(self.flatHyperParameters)
        self.hyperParameterNames.append(self.flatHyperParameterNames[:])
        self.hyperGridConstants.append(self.hyperGridConstant[:])
        self.hyperPrior.append(self.flatHyperPriors[:])
        self.hyperPriorValues.append(self.flatHyperPriorValues[:])

        # count individual transition models
        self.tmCounts = []
        for hpv in self.hyperParameterValues:
            if len(hpv) > 0:
                self.tmCounts.append(len(hpv))
            else:
                self.tmCounts.append(1)
        self.tmCount = np.sum(self.tmCounts)

        if len(self.hyperGridValues) > 0:
            print('+ Added transition model: {} ({} combination(s) of the following hyper-parameters: {})'
                  .format(name, len(self.hyperGridValues), self.hyperParameterNames[-1]))
        else:
            print('+ Added transition model: {} (no hyper-parameters)'.format(name))

    def addTM(self, name, transitionModel):
        """
        See :meth:`.OnlineStudy.addTransitionModel`.
        """
        self.addTransitionModel(name, transitionModel)

    def add(self, name, transitionModel):
        """
        See :meth:`.OnlineStudy.addTransitionModel`.
        """
        self.addTransitionModel(name, transitionModel)

    def setTransitionModelPrior(self, transitionModelPrior, silent=False):
        """
        Sets prior probabilities for transition models added to the online study instance.

        Args:
            transitionModelPrior: List/Array of probabilities, one for each transition model. If the list does not sum
                to one, it will be re-normalised.
            silent: If true, no output is generated by this method.
        """
        if not (isinstance(transitionModelPrior, Iterable) and len(transitionModelPrior) == len(self.transitionModels)):
            raise ConfigurationError('Length of transition model prior ({}) does not fit number of transition models '
                                     '({})'.format(len(transitionModelPrior), len(self.transitionModels)))

        self.transitionModelPrior = np.array(transitionModelPrior)

        if not np.sum(transitionModelPrior) == 1.:
            print('+ WARNING: Transition model prior does not sum up to one. Will re-normalize.')
            self.transitionModelPrior /= np.sum(self.transitionModelPrior)

        if not silent:
            print('+ Set custom transition model prior.')

    def step(self, dataPoint):
        """
        Update the current parameter distribution by adding a new data point to the data set.

        Args:
            dataPoint(float, int, ndarray): Float, int, or 1D-array of those (for multidimensional data).
        """
        # at least one transition model has to be set or added
        if (self.tmCount is None) and (self.transitionModel is None):
            raise ConfigurationError('No transition model set or added.')

        # if one only sets a transition model, but does not use addTransitionModel, we add it here
        if (self.tmCount is None) and (self.transitionModel is not None):
            self.addTransitionModel('transition model', self.transitionModel)

        if not isinstance(dataPoint, list):
            dataPoint = [dataPoint]

        if len(self.rawData) == 0:
            print('+ Start model fit')

            # consistency check to detect duplicate hyper-parameter names in different transition models
            allNames = list(flatten(self.hyperParameterNames))
            if len(allNames) != len(np.unique(allNames)):
                raise ConfigurationError('Detected duplicate hyper-parameter names. Choose unique identifiers.')

            # check the model consistency the first time that 'step' is called
            self.rawData = np.array(dataPoint)
            Study._checkConsistency(self)

            self.rawTimestamps = np.array([0])
            self.formattedTimestamps = []
        else:
            self.rawData = np.append(self.rawData, np.array(dataPoint), axis=0)
            self.rawTimestamps = np.append(self.rawTimestamps, self.rawTimestamps[-1]+1)

        # only proceed if at least one data segment can be created
        if len(self.rawData) < self.observationModel.segmentLength:
            print('+ Not enough data points to start analysis. Will wait for more data.')
            return

        self.formattedTimestamps.append(self.rawTimestamps[-1])

        if self.firstStep:
            # initialize parameter prior distribution
            prior = self._computePrior(silent=False)

            # initialize hyper-prior as flat
            if self.hyperPrior is None:
                self.hyperPrior = 'flat hyper-prior'
                self.hyperPriorValues = [np.ones(tmc) / (tmc * np.prod(hgc))
                                         for tmc, hgc in zip(self.tmCounts, self.hyperGridConstants)]
                print('    + Set flat hyper-prior.')

            # initialize transition model prior as flat
            if self.transitionModelPrior is None:
                self.transitionModelPrior = np.ones(len(self.transitionModels))/len(self.transitionModels)
                if len(self.transitionModels) > 1:
                    print('    + Set flat transition model prior.')

            # initialize logEvidenceList as an array of zeros and add log-lattice-constant for proper normalization
            self.logEvidenceList = [np.zeros(tmc)+np.log(np.prod(self.latticeConstant)) for tmc in self.tmCounts]
            self.hyperLogEvidenceList = np.array([0. for tmc in self.tmCounts])

            # initialize parameter posterior
            self.parameterPosterior = [np.zeros([tmc] + self.gridSize) for tmc in self.tmCounts]

            # initialize hyper-parameter distribution
            self.hyperParameterDistribution = [np.zeros(tmc) for tmc in self.tmCounts]

            # initialize (local) transition model distribution
            self.transitionModelDistribution = np.zeros(len(self.transitionModels))
            self.localTransitionModelDistribution = np.zeros(len(self.transitionModels))

            # initialize transition model posterior
            self.transitionModelPosterior = np.zeros([len(self.transitionModels)] + self.gridSize)

            # initialize marginalized posterior
            self.marginalizedPosterior = np.zeros(self.gridSize)

        # select data segment
        dataSegment = self.rawData[-self.observationModel.segmentLength:]

        # compute current likelihood only once
        likelihood = self.observationModel.processedPdf(self.grid, dataSegment)

        # loop over all hypotheses/transition models
        for i, (tm, hpv) in enumerate(zip(self.transitionModels, self.hyperParameterValues)):
            self.setTransitionModel(tm, silent=True)  # set current transition model

            # account for transition models with no hyper-parameters
            if len(hpv) == 0:
                hpv = [None]

            # loop over all hyper-parameter values to fit
            for j, x in enumerate(hpv):
                # set current hyper-parameter values
                if x is not None:
                    self._setAllHyperParameters(x)

                # compute alpha_i
                if self.firstStep:  # first time step, so use predefined prior
                    alphai = prior*likelihood
                else:  # in all other time steps transform "old" alpha/posterior
                    alphai = self.transitionModel.computeForwardPrior(self.parameterPosterior[i][j],
                                                                      len(self.formattedData)-1)*likelihood
                ni = np.sum(alphai)

                # update log-evidence list
                self.logEvidenceList[i][j] += np.log(ni)
                self.hyperParameterDistribution[i][j] = self.logEvidenceList[i][j] + np.log(self.hyperPriorValues[i][j])

                # store parameter posterior
                self.parameterPosterior[i][j] = alphai/ni

            # marginalize model evidence wrt hyper-parameters, for all past data points and only for the current one
            oldHyperEvidence = self.hyperLogEvidenceList[i]
            self.hyperLogEvidenceList[i] = logsumexp(self.logEvidenceList[i] + np.log(self.hyperPriorValues[i]))
            self.transitionModelDistribution[i] = self.hyperLogEvidenceList[i]
            self.localTransitionModelDistribution[i] = self.hyperLogEvidenceList[i] - oldHyperEvidence + \
                                                       np.log(self.transitionModelPrior[i])

            # normalize hyper-parameter distribution of current transition model
            self.hyperParameterDistribution[i] -= np.amax(self.hyperParameterDistribution[i])
            self.hyperParameterDistribution[i] = np.exp(self.hyperParameterDistribution[i])
            self.hyperParameterDistribution[i] /= np.sum(self.hyperParameterDistribution[i])
            if len(self.hyperGridConstants[i]) > 0:
                self.hyperParameterDistribution[i] /= np.prod(self.hyperGridConstants[i])

            # compute parameter posterior, marginalized wrt the hyper-parameters of the current transition model
            hpd = deepcopy(self.hyperParameterDistribution[i])
            while len(self.parameterPosterior[i].shape) > len(hpd.shape):
                hpd = np.expand_dims(hpd, axis=-1)
            self.transitionModelPosterior[i] = np.sum(self.parameterPosterior[i] *
                                                      hpd *
                                                      np.prod(self.hyperGridConstants[i]), axis=0)

        # normalize (local) transition model distribution (and transform to linear space)
        self.transitionModelDistribution -= np.amax(self.transitionModelDistribution)
        self.transitionModelDistribution = np.exp(self.transitionModelDistribution)
        self.transitionModelDistribution /= np.sum(self.transitionModelDistribution)

        self.localTransitionModelDistribution -= np.amax(self.localTransitionModelDistribution)
        self.localTransitionModelDistribution = np.exp(self.localTransitionModelDistribution)
        self.localTransitionModelDistribution /= np.sum(self.localTransitionModelDistribution)

        # normalize marginalized parameter posterior
        tmd = deepcopy(self.transitionModelDistribution)
        while len(self.transitionModelPosterior.shape) > len(tmd.shape):
            tmd = np.expand_dims(tmd, axis=-1)
        self.marginalizedPosterior = np.sum(self.transitionModelPosterior * tmd, axis=0)

        # compute log-evidence value marginalized over different transition models
        self.logEvidence = logsumexp(self.hyperLogEvidenceList + np.log(self.transitionModelPrior))

        # store results for future plotting
        if self.storeHistory:
            self.posteriorMeanValues.append(np.array([np.sum(self.marginalizedPosterior*g) for g in self.grid]))
            self.posteriorSequence.append(self.marginalizedPosterior.copy())
            self.hyperParameterSequence.append(deepcopy(self.hyperParameterDistribution))
            self.transitionModelSequence.append(deepcopy(self.transitionModelDistribution))
            self.localTransitionModelSequence.append(deepcopy(self.localTransitionModelDistribution))

        if self.firstStep:
            self.firstStep = False

    def fit(self, *args, **kwargs):
        raise NotImplementedError('OnlineStudy object has no "fit" method. Use "step" instead.')

    def getParameterDistribution(self, t, name, plot=False, density=True, **kwargs):
        """
        Compute the marginal parameter distribution at a given time step. Only available if Online Study is created
        with flag 'storeHistory=True'.

        Args:
            t(int, float): Time step/stamp for which the parameter distribution is evaluated
            name(str): Name of the parameter to display
            plot(bool): If True, a plot of the distribution is created
            density(bool): If true, probability density is plotted; if false, probability values.
            **kwargs: All further keyword-arguments are passed to the plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the parameter values, the second one the corresponding
                probability values
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past parameter distributions, Online Study must be called with flag'
                                      '"storeHistory=True". Use "getCurrentParameterDistribution" instead.')

        # plotting function of Study class can only handle arrays, not lists
        self.formattedTimestamps = np.array(self.formattedTimestamps)
        self.posteriorSequence = np.array(self.posteriorSequence)

        x, p = Study.getParameterDistribution(self, t, name, plot=plot, density=density, **kwargs)

        # re-transform arrays to lists, so online study may continue to append values
        self.formattedTimestamps = list(self.formattedTimestamps)
        self.posteriorSequence = list(self.posteriorSequence)
        return x, p

    def getPD(self, t, name, plot=False, density=True, **kwargs):
        """
        See :meth:`.OnlineStudy.getParameterDistribution`.
        """
        return self.getParameterDistribution(t, name, plot=plot, **kwargs)

    def getCurrentParameterDistribution(self, name, plot=False, density=True, **kwargs):
        """
        Compute the current marginal parameter distribution.

        Args:
            name(str): Name of the parameter to display
            plot(bool): If True, a plot of the distribution is created
            density(bool): If true, probability density is plotted; if false, probability values.
            **kwargs: All further keyword-arguments are passed to the plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the parameter values, the second one the corresponding
                probability values
        """
        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        axesToMarginalize = list(range(len(self.observationModel.parameterNames)))
        try:
            axesToMarginalize.remove(paramIndex)
        except ValueError:
            raise PostProcessingError('Wrong parameter index. Available indices: {}'.format(axesToMarginalize))

        x = self.marginalGrid[paramIndex]
        dx = self.latticeConstant[paramIndex]
        marginalDistribution = np.squeeze(
            np.apply_over_axes(np.sum, self.marginalizedPosterior, axesToMarginalize)).copy()

        if density:
            marginalDistribution /= dx

        if plot:
            plt.fill_between(x, 0, marginalDistribution, **kwargs)
            plt.xlabel(self.observationModel.parameterNames[paramIndex])
            if density:
                plt.ylabel('probability density')
            else:
                plt.ylabel('probability')

        return x, marginalDistribution

    def getCPD(self, name, plot=False, density=True, **kwargs):
        """
        See :meth:`.OnlineStudy.getCurrentParameterDistribution`.
        """
        return self.getCurrentParameterDistribution(name, plot=plot, density=density, **kwargs)

    def getParameterDistributions(self, name, plot=False, density=True, **kwargs):
        """
        Computes the time series of marginal posterior distributions with respect to a given model parameter. Only
        available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): Name of the parameter to display
            plot(bool): If True, a plot of the series of distributions is created (density map)
            density(bool): If true, probability density is plotted; if false, probability values.
            **kwargs: All further keyword-arguments are passed to the plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the parameter values, the second one the sequence of
                corresponding posterior distributions.
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past parameter distributions, Online Study must be called with flag'
                                      '"storeHistory=True". Use "getCurrentParameterDistribution" instead.')

        # plotting function of Study class can only handle arrays, not lists
        self.formattedTimestamps = np.array(self.formattedTimestamps)
        self.posteriorSequence = np.array(self.posteriorSequence)

        x, p = Study.getParameterDistributions(self, name, plot=plot, density=density, **kwargs)

        # re-transform arrays to lists, so online study may continue to append values
        self.formattedTimestamps = list(self.formattedTimestamps)
        self.posteriorSequence = list(self.posteriorSequence)
        return x, p

    def getPDs(self, name, plot=False, density=True, **kwargs):
        """
        See :meth:`.OnlineStudy.getParameterDistributions`.
        """
        return self.getParameterDistributions(name, plot=plot, density=density, **kwargs)

    def plotParameterEvolution(self, name, color='b', gamma=0.5, **kwargs):
        """
        Plots a series of marginal posterior distributions corresponding to a single model parameter, together with the
        posterior mean values. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): Name of the parameter to display
            color: color from which a light colormap is created
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5
            kwargs: all further keyword-arguments are passed to the plot of the posterior mean values
        """
        if not self.storeHistory:
            raise PostProcessingError('To plot past parameter distributions, Online Study must be called with flag'
                                      '"storeHistory=True". Use "getCurrentParameterDistribution" instead.')

        # plotting function of Study class can only handle arrays, not lists
        self.formattedTimestamps = np.array(self.formattedTimestamps)
        self.posteriorMeanValues = np.array(self.posteriorMeanValues).T
        self.posteriorSequence = np.array(self.posteriorSequence)

        Study.plotParameterEvolution(self, name, color=color, gamma=gamma, **kwargs)

        # re-transform arrays to lists, so online study may continue to append values
        self.formattedTimestamps = list(self.formattedTimestamps)
        self.posteriorMeanValues = list(self.posteriorMeanValues.T)
        self.posteriorSequence = list(self.posteriorSequence)

    def getCurrentTransitionModelDistribution(self, local=False):
        """
        Returns the current probabilities for each transition model defined in the Online Study.

        Args:
            local(bool): If true, transition model distribution taking into account only the last data point is returned.

        Returns:
            ndarray, ndarray: Arrays of transition model names and normalized probabilities.
        """
        if local:
            return np.array(self.transitionModelNames), self.localTransitionModelDistribution
        else:
            return np.array(self.transitionModelNames), self.transitionModelDistribution

    def getCTMD(self, local=False):
        """
        See :meth:`.OnlineStudy.getCurrentTransitionModelDistribution`.
        """
        return self.getCurrentTransitionModelDistribution(local=local)

    def getCurrentTransitionModelProbability(self, transitionModel, local=False):
        """
        Returns the current posterior probability for a specified transition model.

        Args:
            transitionModel(str): Name of the transition model
            local(bool): If true, transition model probability taking into account only the last data point is returned.

        Returns:
            float: Posterior probability value for the specified transition model
        """
        transitionModelIndex = self.transitionModelNames.index(transitionModel)
        if local:
            return self.localTransitionModelDistribution[transitionModelIndex]
        else:
            return self.transitionModelDistribution[transitionModelIndex]

    def getCTMP(self, transitionModel, local=False):
        """
        See :meth:`.OnlineStudy.getCurrentTransitionModelProbability.
        """
        return self.getCurrentTransitionModelProbability(transitionModel, local=local)

    def getTransitionModelDistributions(self, local=False):
        """
        The transition model distribution contains posterior probability values for all transition models included in
        the online study. This distribution is available for all time steps analyzed. Only available if Online Study
        is created with flag 'storeHistory=True'.

        Args:
            local(bool): If true, transition model distributions taking into account only the data point at the
                corresponding time step is returned.

        Returns:
            ndarray, ndarray: Arrays containing the names and posterior probability values for all transition models
                included in the online study for all time steps analyzed
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past transition model distributions, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentTransitionModelDistribution" instead.')
        if local:
            return np.array(self.transitionModelNames), np.array(self.localTransitionModelSequence)
        else:
            return np.array(self.transitionModelNames), np.array(self.transitionModelSequence)

    def getTransitionModelProbabilities(self, transitionModel, local=False):
        """
        Returns posterior probability values for a specified transition model. This distribution is available for all
        time steps analyzed. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            transitionModel(str): Name of the transition model
            local(bool): If true, transition model probabilities taking into account only the data point at the
                corresponding time step is returned.

        Returns:
            ndarray: Array containing the posterior probability values for the specified transition model for all time
                steps analyzed
            transitionModel(str): Name of the transition model
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past transition model distributions, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentTransitionModelDistribution" instead.')

        transitionModelIndex = self.transitionModelNames.index(transitionModel)

        if local:
            return np.array(self.localTransitionModelSequence)[:, transitionModelIndex]
        else:
            return np.array(self.transitionModelSequence)[:, transitionModelIndex]

    def getTMPs(self, transitionModel, local=False):
        """
        See :meth:`.OnlineStudy.getTransitionModelProbabilities.
        """
        return self.getTransitionModelProbabilities(transitionModel, local=local)

    def getCurrentParameterMeanValue(self, name):
        """
        Returns the posterior mean value for a given parameter of the observation model.

        Args:
            name(str): Name of the parameter

        Returns:
            float: posterior mean value
        """
        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        mean = np.sum(self.marginalizedPosterior*self.grid[paramIndex])
        return mean

    def getParameterMeanValue(self, t, name):
        """
        Returns the posterior mean value for a given parameter of the observation model at a specified time step. Only
        available if Online Study is created with flag 'storeHistory=True'.

        Args:
            t(int): Time step at which to compute parameter mean value
            name(str): Name of the parameter

        Returns:
            float: posterior mean value
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past parameter mean values, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentParameterMeanValue" instead.')

        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        # access parameter distribution
        if t not in self.formattedTimestamps:
            raise PostProcessingError('Supplied time ({}) does not exist in data or is out of range.'.format(t))
        timeIndex = list(self.formattedTimestamps).index(t)  # to select corresponding posterior distribution

        parameterDistribution = self.posteriorSequence[timeIndex]

        mean = np.sum(parameterDistribution[t] * self.grid[paramIndex])
        return mean

    def getParameterMeanValues(self, name):
        """
        Returns the posterior mean value for a given parameter of the observation model for all time steps. Only
        available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): Name of the parameter

        Returns:
            ndarray: posterior mean values
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past parameter mean values, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentParameterMeanValue" instead.')

        # get parameter index
        paramIndex = -1
        for i, n in enumerate(self.observationModel.parameterNames):
            if n == name:
                paramIndex = i

        # check if match was found
        if paramIndex == -1:
            raise PostProcessingError('Wrong parameter name. Available options: {0}'
                                      .format(self.observationModel.parameterNames))

        mean = np.array(self.posteriorMeanValues).T[paramIndex]
        return mean

    def getHyperParameterMeanValue(self, t, name):
        """
        Computes the mean value of the joint hyper-parameter distribution for a given hyper-parameter at a
        given time step. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            t(int): Time step at which to compute distribution
            name(str): name of hyper-parameter

        Returns:
            ndarray: Array containing the mean values of all hyper-parameters of the given transition model
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past hyper-parameter mean values, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentHyperParameterMeanValue" instead.')

        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        # access hyper-parameter distribution
        if t not in self.formattedTimestamps:
            raise PostProcessingError('Supplied time ({}) does not exist in data or is out of range.'.format(t))
        timeIndex = list(self.formattedTimestamps).index(t)  # to select corresponding posterior distribution

        hyperParameterDistribution = self.hyperParameterSequence[timeIndex]

        hyperParameterDistribution = hyperParameterDistribution[tmIndex][:, None]
        hyperParameterValues = self.hyperParameterValues[tmIndex]
        hyperGridConstants = self.hyperGridConstants[tmIndex]

        # compute mean value
        mean = np.sum(hyperParameterValues*hyperParameterDistribution*np.prod(hyperGridConstants), axis=0)
        return mean[hpIndex]

    def getHyperParameterMeanValues(self, name):
        """
        Computes the sequence of mean value of the joint hyper-parameter distribution for a given hyper-parameter for
        all time steps. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): name of hyper-parameter

        Returns:
            ndarray: Array containing the sequences of mean values of the given hyper-parameter
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past hyper-parameter mean values, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentHyperParameterMeanValue" instead.')

        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        hyperParameterSequence = np.array([hp[tmIndex].tolist()
                                           for hp in self.hyperParameterSequence])[:, :, None]
        hyperParameterValues = self.hyperParameterValues[tmIndex]
        hyperGridConstants = self.hyperGridConstants[tmIndex]

        # compute mean value
        mean = np.sum(hyperParameterSequence * hyperParameterValues * np.prod(hyperGridConstants), axis=1).T
        return mean[hpIndex]

    def getHyperParameterDistribution(self, t, name, plot=False, **kwargs):
        """
        Computes marginal hyper-parameter distribution of a single hyper-parameter at a specific time step in an
        OnlineStudy fit. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            t(int): Time step at which to compute distribution, or 'avg' for time-averaged distribution
            name(str): hyper-parameter name
            plot(bool): If True, a bar chart of the distribution is created
            **kwargs: All further keyword-arguments are passed to the bar-plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the hyper-parameter values, the second one the
                corresponding probability values
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past hyper-parameter distributions, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentHyperParameterDistribution" instead.')

        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        # access hyper-parameter distribution
        if t == 'avg':
            # compute time-averaged distribution
            hyperParameterDistribution = np.sum(self.hyperParameterSequence, axis=0)/len(self.hyperParameterSequence)
        else:
            # try to access distribution of specified time step
            if t not in self.formattedTimestamps:
                raise PostProcessingError('Supplied time ({}) does not exist in data or is out of range.'.format(t))
            timeIndex = list(self.formattedTimestamps).index(t)  # to select corresponding posterior distribution

            hyperParameterDistribution = self.hyperParameterSequence[timeIndex]

        hyperParameterDistribution = hyperParameterDistribution[tmIndex]
        axesToMarginalize = list(range(len(self.hyperParameterNames[tmIndex])))
        axesToMarginalize.remove(hpIndex)

        # reshape hyper-parameter grid for easy marginalization
        hyperGridSteps = [len(x) for x in self.allFlatHyperParameterValues[tmIndex]]
        distribution = hyperParameterDistribution.reshape(hyperGridSteps, order='C')
        marginalDistribution = np.squeeze(np.apply_over_axes(np.sum, distribution, axesToMarginalize))

        x = self.allFlatHyperParameterValues[tmIndex][hpIndex]
        if plot:
            # check if categorical
            if np.any(np.abs(np.diff(np.diff(x))) > 10 ** -10):
                plt.bar(np.arange(len(x)), marginalDistribution, align='center', width=1., **kwargs)
                plt.xticks(np.arange(len(x)), x)
                plt.ylabel('probability')
            # regular spacing
            else:
                plt.bar(x, marginalDistribution, align='center',
                        width=self.hyperGridConstants[tmIndex][hpIndex],
                        **kwargs)
                plt.ylabel('probability')

            plt.xlabel(self.hyperParameterNames[tmIndex][hpIndex])

        return x, marginalDistribution

    def getHPD(self, t, name, plot=False, **kwargs):
        """
        See :meth:`.OnlineStudy.getHyperParameterDistribution.
        """
        return self.getHyperParameterDistribution(t, name, plot=plot, **kwargs)

    def getCurrentHyperParameterDistribution(self, name, plot=False, **kwargs):
        """
        Computes marginal hyper-parameter distribution of a single hyper-parameter at the current time step in an
        OnlineStudy fit.

        Args:
            name(str): hyper-parameter name
            plot(bool): If True, a bar chart of the distribution is created
            **kwargs: All further keyword-arguments are passed to the bar-plot (see matplotlib documentation)

        Returns:
            ndarray, ndarray: The first array contains the hyper-parameter values, the second one the
                corresponding probability values
        """
        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        hyperParameterDistribution = self.hyperParameterDistribution[tmIndex]
        axesToMarginalize = list(range(len(self.hyperParameterNames[tmIndex])))
        axesToMarginalize.remove(hpIndex)

        # reshape hyper-parameter grid for easy marginalization
        hyperGridSteps = [len(x) for x in self.allFlatHyperParameterValues[tmIndex]]
        distribution = hyperParameterDistribution.reshape(hyperGridSteps, order='C')
        marginalDistribution = np.squeeze(np.apply_over_axes(np.sum, distribution, axesToMarginalize))
        marginalDistribution *= np.prod(self.hyperGridConstants[tmIndex])

        x = self.allFlatHyperParameterValues[tmIndex][hpIndex]
        if plot:
            # check if categorical
            if np.any(np.abs(np.diff(np.diff(x))) > 10 ** -10):
                plt.bar(np.arange(len(x)), marginalDistribution, align='center', width=1., **kwargs)
                plt.xticks(np.arange(len(x)), x)
                plt.ylabel('probability')
            # regular spacing
            else:
                plt.bar(x, marginalDistribution, align='center',
                        width=self.hyperGridConstants[tmIndex][hpIndex],
                        **kwargs)
                plt.ylabel('probability')

            plt.xlabel(self.hyperParameterNames[tmIndex][hpIndex])

        return x, marginalDistribution

    def getCHPD(self, name, plot=False, **kwargs):
        """
        See :meth:`.OnlineStudy.getCurrentHyperParameterDistribution.
        """
        return self.getCurrentHyperParameterDistribution(name, plot=plot, **kwargs)

    def getHyperParameterDistributions(self, name):
        """
        Computes marginal hyper-parameter distributions of a single hyper-parameter for all time steps in an
        OnlineStudy fit. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): hyper-parameter name

        Returns:
            ndarray, ndarray: The first array contains the hyper-parameter values, the second one the
                corresponding probability values (first axis is time).
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past hyper-parameter distributions, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentHyperParameterDistributions instead.')

        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        hyperParameterSequence = np.array([x[tmIndex] for x in self.hyperParameterSequence])

        # marginalize the hyper-posterior probabilities
        hpv = np.array(self.hyperParameterValues[tmIndex])
        paramHpv = hpv[:, hpIndex]
        uniqueValues = np.sort(np.unique(paramHpv))

        marginalDistribution = []
        for value in uniqueValues:
            marginalDistribution.append([])
            indices = np.where(paramHpv == value)
            for hp in hyperParameterSequence:
                probabilities = hp[indices]
                marginalDistribution[-1].append(np.sum(probabilities))
        marginalDistribution = np.array(marginalDistribution).T

        # renormalize marginal probability distribution
        marginalDistribution /= np.sum(marginalDistribution, axis=1)[:, None]

        return uniqueValues, marginalDistribution

    def getHPDs(self, name):
        """
        See :meth:`.OnlineStudy.getHyperParameterDistributions.
        """
        return self.getHyperParameterDistributions(name)

    def plotHyperParameterEvolution(self, name, color='b', gamma=0.5, **kwargs):
        """
        Plot method to display a series of marginal posterior distributions corresponding to a single model parameter.
        This method includes the removal of plotting artefacts, gamma correction as well as an overlay of the posterior
        mean values. Only available if Online Study is created with flag 'storeHistory=True'.

        Args:
            name(str): hyper-parameter name
            color: color from which a light colormap is created
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5
            kwargs: all further keyword-arguments are passed to the plot of the posterior mean values
        """
        if not self.storeHistory:
            raise PostProcessingError('To get past hyper-parameter distributions, Online Study must be called with '
                                      'flag "storeHistory=True". Use "getCurrentHyperParameterDistribution" instead.')

        # determine indices of transition model and hyper-parameter
        hpIndex = -1
        for i, tm in enumerate(self.transitionModels):
            try:
                hpIndex = self._getHyperParameterIndex(tm, name)
                tmIndex = i
            except PostProcessingError:
                pass
        if hpIndex == -1:
            raise PostProcessingError('No hyper-parameter "{}" found. Check hyper-parameter names.'.format(name))

        # get sequence of hyper-parameter distributions
        uniqueValues, marginalDistribution = self.getHyperParameterDistributions(name)

        # compute hyper-posterior mean values
        meanValues = self.getHyperParameterMeanValues(name)

        # clean up very small probability values, as they may create image artefacts
        pmax = np.amax(marginalDistribution)
        marginalDistribution[marginalDistribution < pmax * (10 ** -20)] = 0

        plt.imshow(marginalDistribution.T ** gamma,
                   origin='lower',
                   cmap=createColormap(color),
                   extent=[self.formattedTimestamps[0], self.formattedTimestamps[-1]] +
                          [uniqueValues[0], uniqueValues[-1]],
                   aspect='auto')

        # set default color of plot to black
        if ('c' not in kwargs) and ('color' not in kwargs):
            kwargs['c'] = 'k'

        # set default linewidth to 1.5
        if ('lw' not in kwargs) and ('linewidth' not in kwargs):
            kwargs['lw'] = 1.5

        plt.plot(self.formattedTimestamps, meanValues, **kwargs)

        plt.ylim(uniqueValues[0], uniqueValues[-1])
        plt.ylabel(self.hyperParameterNames[tmIndex][hpIndex])
        plt.xlabel('time step')

    def getJointHyperParameterDistribution(self, names, plot=False, figure=None, subplot=111, **kwargs):
        raise NotImplementedError('This method is not available in "OnlineStudy".')

    def plot(self, name, **kwargs):
        """
        Convenience method to plot the temporal evolution of (hyper-)parameters, the distribution of a
        (hyper-)parameter at a specific time step, or the temporal evolution of the probability of a transition model.

        Args:
            name(str): name of the (hyper-)parameter or transition model to display
            color: color from which a light colormap is created (for (hyper-)parameter evolution only)
            gamma(float): exponent for gamma correction of the displayed marginal distribution; default: 0.5 (for
                (hyper-)parameter evolution only)
            t: Time step/stamp for which the parameter distribution is evaluated
            density(bool): If true, probability density is plotted; if false, probability values. Note: Only availble
                for parameters, not hyper-parameters.
            local(bool): If true, transition model probabilities taking into account only the data point at the
                corresponding time step is returned
            kwargs: all further keyword-arguments are passed to the axes object of the plot
        """
        density = kwargs.pop('density', True)

        # check if time-step/stamp is supplied
        t = kwargs.pop('t', None)

        # check if results were stored
        if t is not None and not self.storeHistory:
            raise PostProcessingError('Online study has only stored current parameter data ("storeHistory=False"), '
                                      'no time step can be specified, only current (hyper-)parameter distributions will'
                                      'be plotted.')

        # check if name belongs to hyper-parameter
        hyper = False
        for tm in self.transitionModels:
            try:
                self._getHyperParameterIndex(tm, name)
                hyper = True
            except PostProcessingError:
                pass

        # check if name belongs to transition model
        try:
            transitionModelIndex = self.transitionModelNames.index(name)
            model = True
        except ValueError:
            model = False

        if not hyper and not model:
            # parameter evolution plot
            if t is None and self.storeHistory:
                # read additional kwargs (or set default values)
                color = kwargs.pop('color', 'b')
                gamma = kwargs.pop('gamma', 0.5)
                self.plotParameterEvolution(name, color=color, gamma=gamma, **kwargs)

            # parameter distribution plot
            elif t is not None and self.storeHistory:
                self.getParameterDistribution(t, name, plot=True, density=density, **kwargs)
            else:
                self.getCurrentParameterDistribution(name, plot=True, density=density, **kwargs)

        elif hyper and not model:
            # hyper-parameter evolution plot
            if t is None and self.storeHistory:
                # read additional kwargs (or set default values)
                color = kwargs.pop('color', 'b')
                gamma = kwargs.pop('gamma', 0.5)
                self.plotHyperParameterEvolution(name, color=color, gamma=gamma, **kwargs)

            # hyper-parameter distribution plot
            elif t is not None and self.storeHistory:
                self.getHyperParameterDistribution(t, name, plot=True, **kwargs)
            else:
                self.getCurrentHyperParameterDistribution(name, plot=True, **kwargs)

        elif model and not hyper:
            # read additional kwargs (or set default values)
            local = kwargs.pop('local', False)

            # plot local transition model probabilities
            if local:
                plt.plot(self.formattedTimestamps,
                         np.array(self.localTransitionModelSequence)[:, transitionModelIndex],
                         **kwargs)

            # plot transition model probabilities
            else:
                plt.plot(self.formattedTimestamps,
                         np.array(self.transitionModelSequence)[:, transitionModelIndex],
                         **kwargs)

        else:
            raise PostProcessingError('Duplicate names of hyper-/parameters/transition models, cannot use "plot" '
                                      'method.')
