function [MDP] = spm_MDP_VB_XXX(MDP,OPTIONS)
% active inference and learning using belief propagation (factorised)
% FORMAT [MDP] = spm_MDP_VB_XXX(MDP,OPTIONS)
%
% Input; MDP(m,n)       - structure array of m models over n epochs
% MDP.U(1,F)            - controllable factors
% MDP.T                 - number of outcomes
%
% MDP.A{G}(O,N1,...,NF) - likelihood of O outcomes for modality G, given hidden states
% MDP.B{F}(N,N,U)       - transitions among N states under U control states
% MDP.C{G}(O)           - prior probabilities over final outcomes (log preferences)
% MDP.D{F}(N,1)         - prior probabilities over initial states (Dirichlet counts)
% MDP.E(P,1)            - prior probabilities over control states (Dirichlet counts)
%
% MDP.a{G}              - concentration parameters for A
% MDP.b{F}              - concentration parameters for B
% MDP.c{G}              - concentration parameters for C
% MDP.d{F}              - concentration parameters for D
% MDP.e{P}              - concentration parameters for E
%
% optional:
% MDP.s(F,T)            - true states   - for each hidden factor
% MDP.o(G,T)            - true outcomes - for each outcome modality
% MDP.O{G,T}            - likelihoods   - for each outcome modality
% MDP.u(F,T)            - true controls - for each hidden factor
%
% MDP.alpha             - precision - action selection [512]
% MDP.chi               - Occams window for deep updates
% MDP.eta               - Forgetting hyperparameter [1]
% MDP.N                 - depth of deep policy search [N = 0]
% MDP.k(1,F)            - beliefs about controllable factors
%
% MDP.demi.C            - Mixed model: cell array of true causes (DEM.C)
% MDP.demi.U            - Bayesian model average (DEM.U) see: spm_MDP_DEM
% MDP.link              - link array to generate outcomes from
%                         subordinate MDP for deep (hierarchical) models
%
% MDP.n(O,T)            - outputs for modality O at time T are generated by
%                         agent n(O,T); unless n(O,T) = 0, when outputs
%                         are generated by the agents states
% MDP.m(F)              - states for factor F are generated for agent m(F);
%                         unless m(F) = 0, when states are updated for the
%                         agent in question
%
% OPTIONS.A             - switch to evaluate explicit action
% OPTIONS.B             - switch to evaluate backwards pass (replay)
% OPTIONS.N             - switch to evaluate neuronal responses
% OPTIONS.P             - switch to plot graphics: [default: 0)
% OPTIONS.D             - switch to update initial states with final states
% OPTIONS.BMR           - Bayesian model reduction for multiple trials
%                         see: spm_MDP_VB_sleep(MDP,BMR)
% Outputs:
%
% MDP.P{F}(U,T)         - conditional expectations over control states
% MDP.X{F}(N,T)         - conditional expectations over hidden states
% MDP.Y{O,T}            - conditional expectations over outcomes
% MDP.R(P,T)            - conditional expectations over policies
%
% MDP.F(1,T)            - (negative) free energies (states)  over time
% MDP.Z{U,T}            - (negative) free energies (control) over time
% MDP.G{P,T}            - (negative) expected free energies  over time
% MDP.Fa(G)             - (negative) free energy of parameters (a)
% MDP.Fb(F)             - ...
%
% MDP.v                 - expected free energy over policies
% MDP.w                 - precision of beliefs about policies
% MDP.un                - simulated neuronal encoding of hidden states
% MDP.xn                - simulated neuronal encoding of policies
% MDP.wn                - simulated neuronal encoding of precision (tonic)
% MDP.dn                - simulated dopamine responses (phasic)
%
% This routine provides solutions of active inference (minimisation of
% variational free energy) using a generative model based upon a Markov
% decision process. The model and inference scheme is formulated in
% discrete space and time. This means that the generative model (and
% process) are hidden Markov models whose dynamics are given by transition
% probabilities among states and the likelihood corresponds to a particular
% outcome conditioned upon hidden states.
%
% This implementation equips agents with the prior beliefs that they will
% maximise expected free energy. Expected free energy can be interpreted in
% several ways - most intuitively as minimising the KL divergence between
% predicted and preferred outcomes (specified as prior beliefs) -i.e., risk
% while simultaneously minimising ambiguity. Alternatively, this can be
% rearranged into expected information gain and expected value, where value
% is the log of prior preferences (overstates or outcomes).
%
% This implementation generalises previous MDP based formulations of active
% inference by equipping each factor of latent states with a number of
% paths; some of which may be controllable and others not. Controllable
% factors are now specified with indicator variables in the vector MDP.U.
% Furthermore, because the scheme uses sophisticated inference (i.e., a
% recursive tree search accumulating path integral is of expected free
% energy) a policy reduces to a particular combination of controllable
% paths or dynamics over factors. In consequence, posterior beliefs cover
% latent states and paths; with their associated variational free energies.
% Furthermore, it is now necessary to specify the initial states and the
% initial paths using D and E respectively. In other words, he now plays
% the role of a prior over the path of each factor that can only be changed
% if it is controllable (it no longer corresponds to a prior over
% policies).
%
% In addition to state and path estimation (and policy selection), the
% scheme also updates model parameters; including the state transition
% matrices, mapping to outcomes and the initial state. This is useful for
% learning the context. Likelihood and prior probabilities can be specified
% in terms of concentration parameters (of a Dirichlet distribution
% (a,b,c,...). If the corresponding (A,B,C,...) are supplied, they will be
% used to generate outcomes.
%
% This scheme allows for differences in the functional form of priors –
% specified in terms of probability transition tensors – between the
% generating process and generative model. The generative model is, by
% default, specified in terms of Dirichlet parameters, while the generative
% process is specified in terms of expected (likelihood and prior
% transition) probabilities: b and B, respectively. If the number or
% dimensionality of b and B do not correspond, then select OPTIONS.A = 1.
% This will automatically evaluate the most likely policy (combination of
% controllable paths) to reproduce the predicted outcomes (i.e. that which
% minimises variational free energy or maximises accuracy); as opposed to
% using the path selected by the model.
%
% scheme is designed for any allowable policies or control variables
% specified in MDP.U. Constraints on allowable policies can limit the
% numerics or combinatorics considerably. Further, the outcome space and
% hidden states can be defined in terms of factors; corresponding to
% sensory modalities and (functionally) segregated representations,
% respectively. This means, for each factor or subset of hidden states
% there are corresponding control states that determine the transition
% probabilities. in this implementation, hidden factors are combined using
% a Kronecker intensive product to enable exact Bayesian inference using
% belief propagation (the Kronecker tensor form ensures that conditional
% dependencies among hidden factors are evaluated).
%
% In this belief propagation scheme, the next action is evaluated in terms
% of the free energy expected under all subsequent actions until some time
% horizon (specified by MDP.T). This expected free energy is accumulated
% along all allowable paths or policies (see the subroutine spm_forward);
% effectively, performing a deep tree search over future sequences of
% actions. Because actions are conditionally independent of previous
% actions, it is only necessary to update posterior beliefs over hidden
% states at the current time point (using a Bayesian belief updating) and
% then use the prior over actions (based upon expected free energy) to
% select the next action. Previous actions are inferred under the posterior
% beliefs over current states; i.e., inferred state transitions.
%
% In brief, the agent encodes beliefs about hidden states in the past
% conditioned on realised outcomes and actions. The resulting conditional
% expectations determine the (path integral) of free energy that then
% determines an empirical prior over the next action, from which the next
% realised action sampled
%
%
% If supplied with a structure array, this routine will automatically step
% through the implicit sequence of epochs (implicit in the number of
% columns of the array). If the array has multiple rows, each row will be
% treated as a separate model or agent. This enables agents to communicate
% through acting upon a common set of hidden factors, or indeed sharing the
% same outcomes.
%
% See also: spm_MDP, which uses multiple future states and a mean field
% approximation for control states - but allows for different actions at
% all times (as in control problems).
%
% See also: spm_MDP_VB_X,  which is the corresponding variational message
% passing scheme for fixed policies; i.e., ordered sequences of actions
% that are specified a priori.
%
% See also: spm_MDP_VB_XX,  which is the corresponding variational message
% passing scheme for sophisticated policy searches under the assumption
% that the generative process and model have the same structure
%
%
%__________________________________________________________________________
% Copyright (C) 2019 Wellcome Trust Centre for Neuroimaging

% Karl Friston
% $Id: spm_MDP_VB_XXX.m 8418 2023-02-27 19:17:56Z karl $


% deal with a sequence of trials
%==========================================================================

% options
%--------------------------------------------------------------------------
try, OPTIONS.A; catch, OPTIONS.A = 0; end      % action selection
try, OPTIONS.B; catch, OPTIONS.B = 0; end      % backwards pass
try, OPTIONS.D; catch, OPTIONS.D = 0; end      % final states
try, OPTIONS.N; catch, OPTIONS.N = 0; end      % neuronal responses
try, OPTIONS.O; catch, OPTIONS.O = 0; end      % probabilistic outcomes
try, OPTIONS.P; catch, OPTIONS.P = 0; end      % graphics


% check MDP specification
%--------------------------------------------------------------------------
MDP = spm_MDP_check(MDP);

% handle multiple trials, ensuring parameters (and posteriors) are updated
%==========================================================================
if size(MDP,2) > 1

    % plotting options
    %----------------------------------------------------------------------
    GRAPH     = OPTIONS.P;
    OPTIONS.P = 0;

    for i = 1:size(MDP,2)                      % number of trials
        for m = 1:size(MDP,1)                  % number of agents
            if i > 1                           % if previous inversions

                % update concentration parameters
                %----------------------------------------------------------
                MDP(m,i) = spm_MDP_update(MDP(m,i),OUT(m,i - 1));

                % update initial states (post-diction)
                %----------------------------------------------------------
                if any(OPTIONS.D)
                    nD = numel(MDP(m,i).D);
                    if numel(OPTIONS.D) ~= nD
                        OPTIONS.D = ones(nD,1);
                    end
                    for f = 1:nD
                        if OPTIONS.D(f)
                            MDP(m,i).D{f} = OUT(m,i - 1).X{f}(:,end);
                        end
                    end
                end
            end
        end

        % solve this trial (for all models synchronously)
        %------------------------------------------------------------------
        OUT(:,i) = spm_MDP_VB_XXX(MDP(:,i),OPTIONS);
        fprintf('trial %i\n',i);

        % Bayesian model reduction
        %------------------------------------------------------------------
        if isfield(OPTIONS,'BMR')
            if isfield(OPTIONS.BMR,'fun')
                bmrfun = OPTIONS.BMR.fun;
            else
                bmrfun = @(MDP,BMR) spm_MDP_VB_sleep(MDP,BMR);
            end
            for m = 1:size(MDP,1)
                OUT(m,i) = bmrfun(OUT(m,i),OPTIONS.BMR);
            end
        end

    end

    % plot summary statistics - over trials
    %----------------------------------------------------------------------
    MDP = OUT;
    if GRAPH
        if ishandle(GRAPH)
            figure(GRAPH); clf
        else
            spm_figure('GetWin','MDP'); clf
        end
        spm_MDP_VB_game(MDP(1,:))
    end
    return
end


% set up and preliminaries
%==========================================================================

% defaults
%--------------------------------------------------------------------------
try, alpha = MDP(1).alpha; catch, alpha = 512;  end % action precision
try, eta   = MDP(1).eta;   catch, eta   = 1;    end % learning rate
try, chi   = MDP(1).chi;   catch, chi   = 1/64; end % Occam window updates
try, N     = MDP(1).N;     catch, N     = 0;    end % depth of policy search


% initialise model-specific parameters
%==========================================================================
for m = 1:numel(MDP)

    % number of outcomes, states, controls and policies
    %----------------------------------------------------------------------
    Ng(m) = numel(MDP(m).A);                   % number of outcome factors
    Nf(m) = numel(MDP(m).B);                   % number of hidden factors
    for g = 1:Ng(m)
        No(m,g) = size(MDP(m).A{g},1);         % number of outcomes
    end
    for f = 1:Nf(m)
        Ns(m,f) = size(MDP(m).B{f},1);         % number of hidden states
        Nu(m,f) = size(MDP(m).B{f},3);         % number of hidden controls
    end

    % allow for differences between the model (b) and process (B)
    %----------------------------------------------------------------------
    NF    = Nf;
    NS    = Ns;
    NU    = Nu;
    if isfield(MDP(m),'b')
        Nf(m) = numel(MDP(m).b);               % number of hidden factors
        for f = 1:Nf(m)
            Ns(m,f) = size(MDP(m).b{f},1);     % number of hidden states
            Nu(m,f) = size(MDP(m).b{f},3);     % number of hidden controls
        end
    end

    % parameters of generative model and policies
    %======================================================================

    % likelihood model (for a partially observed MDP)
    %----------------------------------------------------------------------
    for g = 1:Ng(m)

        % parameters (concentration parameters): a
        %------------------------------------------------------------------
        if isfield(MDP,'a')
            a{m,g} = MDP(m).a{g};
        else
            a{m,g} = MDP(m).A{g}*512;
        end

        % normalised likelihood
        %------------------------------------------------------------------
        A{m,g} = spm_norm(a{m,g});

        % prior concentration parameters and novelty (W)
        %------------------------------------------------------------------
        if isfield(MDP,'a')
            W{m,g} = spm_wnorm(a{m,g}).*(a{m,g} > 0);
        else
            W{m,g} = false;
        end

        % and ambiguity (H) (for computation of expected free energy: G)
        %------------------------------------------------------------------
        if islogical(a{m,g})
            H{m,g} = false;
        else
            H{m,g} = sum(A{m,g}.*spm_psi(a{m,g}),1);

            if ~any(H{m,g},'all'), H{m,g} = false;  end
        end
    end

    % transition probabilities (priors)
    %----------------------------------------------------------------------
    for f = 1:Nf(m)

        % parameters (concentration parameters): b
        %--------------------------------------------------------------
        if isfield(MDP,'b')
            b{m,f} = MDP(m).b{f};
        else
            b{m,f} = MDP(m).B{f}*512;
        end

        % normalised transition probabilities
        %------------------------------------------------------------------
        B{m,f}     = spm_norm(b{m,f});

        % priors over initial hidden states: concentration parameters
        %------------------------------------------------------------------
        if isfield(MDP,'D')
            D{m,f} = spm_norm(MDP(m).D{f});
        elseif isfield(MDP,'d')
            D{m,f} = spm_norm(MDP(m).d{f});
        else
            D{m,f} = spm_norm(ones(Ns(m,f),1));
        end

        % priors over control states: concentration parameters
        %------------------------------------------------------------------
        if isfield(MDP,'E')
            E{m,f} = spm_norm(MDP(m).E{f});
        elseif isfield(MDP,'e')
            E{m,f} = spm_norm(MDP(m).e{f});
        else
            E{m,f} = spm_norm(ones(Nu(m,f),1));
        end

    end


    % prior preferences (log probabilities) : C
    %----------------------------------------------------------------------
    for g = 1:Ng(m)

        if isfield(MDP,'C')
            C{m,g} = MDP(m).C{g}(:,1);
        elseif isfield(MDP,'c')
            C{m,g} = spm_psi(MDP(m).c{g}(:,1) + 1/32);
        else
            C{m,g} = zeros(No(m,g),1);
        end

        % prior preferences (log probabilities) : C
        %------------------------------------------------------------------
        C{m,g} = spm_log(spm_softmax(C{m,g}));

    end

    % allowable combinations of actions (policies)
    %----------------------------------------------------------------------
    T(m)      = MDP(m).T;                       % number of updates
    U{m}      = any(MDP(m).U,1);                % controllable factors

    % actual policies (J)
    %----------------------------------------------------------------------
    if OPTIONS.A
        k         = find(U{m});
        u         = spm_combinations(NU(m,k));
        J{m}      = zeros(size(u,1),NF(m));
        J{m}(:,k) = u;
    end

    % controllable factors
    %----------------------------------------------------------------------
    if isfield(MDP(m),'k')
        U{m} = any(MDP(m).k,1);
    end

    % beliefs about policies (U)
    %----------------------------------------------------------------------
    if size(MDP(m).U,1) == 1
        k         = find(U{m});
        u         = spm_combinations(Nu(m,k));
        V{m}      = zeros(size(u,1),Nf(m));
        V{m}(:,k) = u;
    else
        V{m} = MDP(m).U;
    end
    Np(m)    = size(V{m},1);                  % number of policies

    
    % initialise posterior expectations (Q) of hidden states at the current
    % time (X) and over time (S)
    %======================================================================
    for f = 1:Nf(m)
        for t = 1:T
            Q{m,f,t} = D{m,f};
        end
        X{m,f} = repmat(D{m,f},1,T);
        S{m,f} = repmat(E{m,f},1,T);

        if OPTIONS.N
            sn{m,f} = zeros(Ns(m,f),T,T) + 1/Ns(m,f);
        end
    end

    % initialise posteriors over control states
    %----------------------------------------------------------------------
    for f = 1:Nf(m)
        for t = 1:T
            P{m,f,t} = E{m,f};
        end
    end

    % if states have not been specified, set to 0
    %----------------------------------------------------------------------
    d  = zeros(NF(m),T);
    try
        i    = find(MDP(m).s);
        d(i) = MDP(m).s(i);
    end
    MDP(m).s = d;

    % if controls have not been specified, set to 0
    %----------------------------------------------------------------------
    d  = zeros(NF(m),T);
    try
        i    = find(MDP(m).u);
        d(i) = MDP(m).u(i);
    end
    MDP(m).u = d;

    % if outcomes have not been specified set to 0
    %======================================================================
    d  = zeros(Ng(m),T);
    try
        i    = find(MDP(m).o);
        d(i) = MDP(m).o(i);
    end
    MDP(m).o = d;

    % sample outcomes if specified probabilistically
    %----------------------------------------------------------------------
    for g  = 1:Ng(m)
        for t = 1:T
            if OPTIONS.O && isfield(MDP(m),'O') 
                O{m,g,t}      = MDP(m).O{g,t};
                po            = O{m,g,t};
                MDP(m).o(g,t) = find(rand < cumsum(po),1);
            else
                O{m,g,t}      = [];
            end
        end
    end


end % end model (m)


% ensure any outcome generating agent is updated first
%--------------------------------------------------------------------------
T       = max(T);                            % maximum number of updates
N       = min(N,T);                          % depth of policy search
[M,MDP] = spm_MDP_get_M(MDP,T,Ng);           % order of model updating


% belief updating over successive time points
%==========================================================================
for t = 1:T

    % generate hidden states and controls for each agent or model
    %======================================================================
    for m = M(t,:)

        % initialise and propagate control state
        %==================================================================
        for f = 1:Nf(m)
            if ~MDP(m).u(f,t)
                if t > 1
                    MDP(m).u(f,t) = MDP(m).u(f,t - 1);
                else

                    % unless it is the initial control state
                    %------------------------------------------------------
                    if U{m}(f)

                        % if controllable sample from prior beliefs
                        %--------------------------------------------------
                        pu = E{m,f};

                    else

                        % otherwise sample a random state
                        %--------------------------------------------------
                        pu = spm_norm(ones(NU(m,f),1));
                    end

                    MDP(m).u(f,t) = find(rand < cumsum(pu),1);
                end
            end
        end

        % posterior predictive density (i.e., prior) over hidden states
        %==================================================================
        for f = 1:Nf(m)
            if t > 1
                Q{m,f,t} = spm_dot(B{m,f},P(m,f,t - 1))*Q{m,f,t - 1};
            else
                Q{m,f,t} = D{m,f};
            end
        end

        % action generating outcomes (for controllable factors)
        %==================================================================
        if t > 1

            % explicit action (e.g., B and b have different structures)
            %--------------------------------------------------------------
            if OPTIONS.A

                % find actions that minimise prediction error
                %----------------------------------------------------------
                i = any(J{m},1);
                for k = 1:size(J{m},1)

                    % predicted states under this policy
                    %------------------------------------------------------
                    u     = MDP(m).u(:,t - 1);
                    u(i)  = J{m}(k,i);
                    for f = 1:numel(MDP(m).B)
                        qs{f} = MDP(m).B{f}(:,MDP(m).s(f,t - 1),u(f));
                    end

                    % free energy or prediction error (i.e., inaccuracy)
                    %------------------------------------------------------
                    F(k)  = 0;
                    for g = 1:Ng(m)
                        po   = spm_dot(MDP(m).A{g},qs);
                        qo   = spm_dot(A{m,g},Q(m,:,t));
                        F(k) = F(k) + qo'*(spm_log(qo) - spm_log(po));
                    end
                end

                % most likely control state
                %----------------------------------------------------------
                [F,k] = min(F);
                MDP(m).u(i,t - 1) = spm_vec(J{m}(k,i));

            else

                % implcit action
                %==========================================================
                for f = 1:NF(m)
                    if U{m}(f)
                        pu                = P{m,f,t - 1};
                        MDP(m).u(f,t - 1) = find(rand < cumsum(pu),1);
                    end
                end

            end

        end % end generating actions


        % sample state if not specified
        %==================================================================
        for f = 1:NF(m)
            if ~MDP(m).s(f,t)

                % the next state is generated by state transititions
                %----------------------------------------------------------
                if t > 1
                    ps = MDP(m).B{f}(:,MDP(m).s(f,t - 1),MDP(m).u(f,t - 1));
                else

                    % unless it is the initial state
                    %------------------------------------------------------
                    if U{m}(f)

                        % if controllable sample from prior beliefs
                        %--------------------------------------------------
                        ps = D{m,f};

                    else

                        % otherwise sample a random state
                        %--------------------------------------------------
                        ps = spm_norm(ones(NS(m,f),1));
                    end
                end

                MDP(m).s(f,t) = find(rand < cumsum(ps),1);

            end

        end % end generating states

    end % end generating over models


    % share states if specified in MDP.m
    %----------------------------------------------------------------------
    for m = M(t,:)
        for f = 1:NF(m)
            if isfield(MDP(m),'m')
                n = MDP(m).m(f);
                if n
                    MDP(m).s(f,t) = MDP(n).s(f,t);
                end
            end
        end
    end

    % generate outcomes for each agent or model
    %======================================================================
    for m = M(t,:)


        % sample outcome, if not specified
        %------------------------------------------------------------------
        for g = 1:Ng(m)

            % if outcome is not specified
            %--------------------------------------------------------------
            if ~MDP(m).o(g,t)

                % outcome is generated by model n
                %----------------------------------------------------------
                if MDP(m).n(g,t) > 0

                    n    = MDP(m).n(g,t);
                    if n == m

                        % outcome that maximises accuracy (i.e. ELBO)
                        %--------------------------------------------------
                        F             = spm_dot(spm_psi(a{m,g}),Q(m,:,t));
                        po            = spm_softmax(F*512);
                        MDP(m).o(g,t) = find(rand < cumsum(po),1);

                    else

                        % outcome from model n (previously sampled)
                        %--------------------------------------------------
                        MDP(m).o(g,t) = MDP(n).o(g,t);

                    end

                elseif MDP(m).n(g,t) < 0

                    % outcome is generated by all models or agents
                    %------------------------------------------------------
                    Fm{g,m} = spm_dot(spm_psi(a{m,g}),Q(m,:,t));


                else

                    % or sample from likelihood, given hidden state
                    %------------------------------------------------------
                    ind           = num2cell(MDP(m).s(:,t));
                    po            = MDP(m).A{g}(:,ind{:});
                    MDP(m).o(g,t) = find(rand < cumsum(po),1);

                end
            end
        end
    end


    % probabilistic outcomes O{m,g,t} from samples or subordinate level
    %======================================================================
    for m = M(t,:)
        for g = 1:Ng(m)

            % share outcomes if specified in MDP.m
            %--------------------------------------------------------------
            if MDP(m).n(g,t) < 0
                F             = sum([Fm{g,:}],2);
                O{m,g,t}      = spm_softmax(F);
                po            = spm_softmax(F*512);
                MDP(m).o(g,t) = find(rand < cumsum(po),1);

            else 

                % use sampled outcomes if no probabilistic specification
                %----------------------------------------------------------
                if isempty(O{m,g,t})
                    O{m,g,t} = full(sparse(MDP(m).o(g,t),1,1,No(m,g),1));
                end

            end
        end
    end


    % or generate outcomes from a subordinate MDP
    %======================================================================
    for m = M(t,:)

        if isfield(MDP,'link')

            % use previous inversions (if available) to generate outcomes
            %--------------------------------------------------------------
            try
                mdp = MDP(m).mdp(t);
            catch
                try
                    mdp = spm_MDP_update(MDP(m).MDP(t),MDP(m).mdp(t - 1));
                catch
                    try
                        mdp = spm_MDP_update(MDP(m).MDP(1),MDP(m).mdp(t - 1));
                    catch
                        mdp = MDP(m).MDP(1);
                    end
                end
            end

            % priors over states (of subordinate level)
            %--------------------------------------------------------------
            mdp.factor = [];
            for f = 1:size(MDP(m).link,1)
                for g = 1:size(MDP(m).link,2)
                    if ~isempty(MDP(m).link{f,g})

                        % subordinate factor has hierarchical constraints
                        %--------------------------------------------------
                        mdp.factor(end + 1) = f;

                        % empirical priors over initial states
                        %--------------------------------------------------
                        O{m,g,t} = spm_dot(A{m,g},Q(m,:,t));
                        mdp.D{f} = MDP(m).link{f,g}*O{m,g,t};

                        % outcomes (i.e., states) are generated by model n
                        %--------------------------------------------------
                        if MDP(m).n(g,t)
                            n    = MDP(m).n(g,t);
                            if m == n
                                ps         = MDP(m).link{f,g}(:,MDP(m).o(g,t));
                                mdp.s(f,1) = find(ps);
                            else
                                mdp.s(f,1) = MDP(n).mdp(t).s(f,1);
                            end
                        end

                        % hidden state for lower level is the outcome
                        %--------------------------------------------------
                        try
                            mdp.s(f,1) = mdp.s(f,1);
                        catch
                            ps         = MDP(m).link{f,g}(:,MDP(m).o(g,t));
                            mdp.s(f,1) = find(ps);
                        end

                    end
                end
            end


            % empirical prior preferences
            %--------------------------------------------------------------
            if isfield(MDP,'linkC')
                for f = 1:size(MDP(m).linkC,1)
                    for g = 1:size(MDP(m).linkC,2)
                        if ~isempty(MDP(m).linkC{f,g})
                            O{m,g,t} = spm_dot(A{m,g},Q(m,:,t));
                            mdp.C{f} = spm_log(MDP(m).linkC{f,g}*O{m,g,t});
                        end
                    end
                end
            end

            % empirical priors over policies
            %--------------------------------------------------------------
            if isfield(MDP,'linkE')
                mdp.factorE = [];
                for g = 1:size(MDP(m).linkE,2)
                    if ~isempty(MDP(m).linkE{g})
                        O{m,g,t} = spm_dot(A{m,g},Q(m,:,t));
                        mdp.E    = MDP(m).linkE{g}*O{m,g,t};
                    end
                end
            end


            % infer hidden states at lower level (outcomes at this level)
            %==============================================================
            OPT.B         = 1;
            MDP(m).mdp(t) = spm_MDP_VB_XXX(mdp,OPT);


            % get inferred outcomes from subordinate MDP
            %==============================================================
            for f = 1:size(MDP(m).link,1)
                for g = 1:size(MDP(m).link,2)
                    if ~isempty(MDP(m).link{f,g})
                        O{m,g,t} = MDP(m).link{f,g}'*MDP(m).mdp(t).X{f}(:,1);
                    end
                end
            end

            % if hierarchical preferences, these contribute to outcomes ...
            %--------------------------------------------------------------
            if isfield(MDP,'linkC')
                for f = 1:size(MDP(m).linkC,1)
                    for g = 1:size(MDP(m).linkC,2)
                        if ~isempty(MDP(m).linkC{f,g})
                            indC     = sparse(MDP(m).mdp(t).o(f,:)',1:length(MDP(m).mdp(t).o(f,:)),ones(length(MDP(m).mdp(t).o(f,:)),1),size(MDP(m).mdp(t).C{f},1),size(MDP(m).mdp(t).C{f},2));
                            O{m,g,t} = spm_softmax(spm_log(O{m,g,t}) + MDP(m).linkC{f,g}'*sum((indC.*(MDP(m).mdp(t).C{f})),2));
                        end
                    end
                end
            end

            % ... and the same for policies
            %--------------------------------------------------------------
            if isfield(MDP,'linkE')
                for g = 1:size(MDP(m).linkE,2)
                    if ~isempty(MDP(m).linkE{g})
                        O{m,g,t} = spm_softmax(spm_log(O{m,g,t}) + spm_log(MDP(m).linkE{g}'*MDP(m).mdp(t).R(:,end)));
                    end
                end
            end

            % ensure DEM starts with final states from previous inversion
            %--------------------------------------------------------------
            if isfield(MDP(m).MDP,'demi')
                MDP(m).MDP.DEM.G(1).x = MDP(m).mdp(t).dem(end).pU.x{1}(:,end);
                MDP(m).MDP.DEM.M(1).x = MDP(m).mdp(t).dem(end).qU.x{1}(:,end);
            end

        end % end of hierarchical mode (link)


        % or generate outcome likelihoods from a variational filter
        %==================================================================
        if isfield(MDP,'demi')

            % use previous inversions (if available)
            %--------------------------------------------------------------
            try
                MDP(m).dem(t) = spm_ADEM_update(MDP(m).dem(t - 1));
            catch
                MDP(m).dem(t) = MDP(m).DEM;
            end

            % get prior over outcomes
            %--------------------------------------------------------------
            for g = 1:Ng(m)
                O{m,g,t} = spm_dot(A{m,g},Q(m,:,t));
            end

            % get posterior outcome from Bayesian filtering
            %--------------------------------------------------------------
            MDP(m).dem(t) = spm_MDP_DEM(MDP(m).dem(t),...
                MDP(m).demi,O(m,:,t),MDP(m).o(:,t));

            for g = 1:Ng(m)
                O{m,g,t} = MDP(m).dem(t).X{g}(:,end);
            end

        end % end outcomes from Bayesian filter

    end % end loop over models or agents


    % Bayesian belief updating hidden states (Q) and controls (P)
    %======================================================================
    for m = M(t,:)

        % belief propagation (K) under policy k
        %------------------------------------------------------------------
        for f = 1:Nf(m)
            for k = 1:Np(m)
                if V{m}(k,f)
                    K{m,f,k} = B{m,f}(:,:,V{m}(k,f));
                else
                    K{m,f,k} = spm_dot(B{m,f},P(m,f,t));
                end
            end
        end

        % posterior over hidden states (Q) and expected free energy (G)
        %==================================================================
        [G,Q,F] = spm_forwards(O,Q,A,K,C,H,W,t,T,min(T,t + N),m);

        % augment with prior probability over paths
        %------------------------------------------------------------------
        for k = 1:Np(m)
            LG     = 0;
            for f  = find(U{m})
                LG = LG + spm_log(E{m,f}(V{m}(k,f)));
            end
            G(k)   = G(k) + LG;
        end

        % prior beliefs about policies (R) and precision (w)
        %------------------------------------------------------------------
        R{m}(:,t)  = spm_softmax(G);
        w{m}(t)    = R{m}(:,t)'*spm_log(R{m}(:,t));
        v{m}(t)    = R{m}(:,t)'*G;


        % posterior over previous control state (P)
        %==================================================================
        Z  = 0;
        if t > 1
            for f = 1:Nf(m)

                % log likelihood of control states
                %----------------------------------------------------------
                LL = spm_dot(spm_dot(spm_psi(b{m,f}),Q{m,f,t}),Q{m,f,t - 1});

                % prior over previous control state
                %----------------------------------------------------------
                LP = spm_log(P{m,f,t - 1});

                % posterior over control states
                %----------------------------------------------------------
                P{m,f,t - 1} = spm_softmax(LL + LP);

                % (negative) complexity of control states
                %----------------------------------------------------------
                Z  = Z + P{m,f,t - 1}'*(LL + LP - spm_log(P{m,f,t - 1}));

            end

        end % end previous control states


        % prior over current control state (P)
        %==================================================================

        % prior over policy (i.e.,action combination)
        %------------------------------------------------------------------
        Ru    = spm_softmax(alpha*G);
        for f = 1:Nf(m)
            
            if U{m}(f)

                % prior over control state
                %----------------------------------------------------------
                for u = 1:Nu(m,f)
                    P{m,f,t}(u) = Ru'*(V{m}(:,f) == u);
                end

            else

                % or previous control state, if path
                %----------------------------------------------------------
                if t > 1
                    P{m,f,t} = P{m,f,t - 1};
                end
            end

        end % end current control states


        % active (likelihood) learning
        %==================================================================

        % mapping from hidden states to outcomes: a
        %------------------------------------------------------------------
        if isfield(MDP(m),'a') && any(U{m})
            for g = 1:Ng(m)

                da     = spm_cross(O(m,g,t),Q{m,:,t});
                da     = reshape(da,[No(m,g),Ns(m,:)]);
                da     = da.*(a{m,g} > 0);

                % update likelihood Dirichlet parameters
                %----------------------------------------------------------
                a{m,g} = (a{m,g} + da)*eta;
                A{m,g} = spm_norm(a{m,g});

                % prior concentration parameters and novelty (W)
                %----------------------------------------------------------
                W{m,g} = spm_wnorm(a{m,g}).*(a{m,g} > 0);

                % and ambiguity (H)
                %----------------------------------------------------------
                H{m,g} = full(sum(A{m,g}.*spm_psi(a{m,g}),1));

            end
        end

        % (ELBO) free energy: states, policies and paths
        %------------------------------------------------------------------
        MDP(m).F(t) = F;
        MDP(m).G{t} = G;
        MDP(m).Z(t) = Z;

        % save marginal posteriors over hidden states: c.f., working memory
        %==================================================================
        if OPTIONS.N
            for f = 1:Nf(m)
                for i = 1:T
                    sn{m,f}(:,i,t) = Q{m,f,i};
                end
            end
        end

        % posterior predictive density Y{g,t}
        %------------------------------------------------------------------
        for g = 1:Ng(m)
            MDP(m).Y{g,t} = spm_dot(A{m,g},Q(m,:,t));
        end

        % return to supraordinate level
        %==================================================================

        % check for residual uncertainty (in hierarchical schemes)
        %------------------------------------------------------------------
        if isfield(MDP,'factor')

            % break if there is no further uncertainty to resolve
            %--------------------------------------------------------------
            for f = MDP(m).factor(:)'
                sq(m,f) = Q{m,f,t}'*spm_log(Q{m,f,t});
            end
            if sum(sq(:)) > - chi
                T = t;
            end
        end


    end % end of loop over models (agents)

    % terminate evidence accumulation
    %----------------------------------------------------------------------
    if t == T
        for m = 1:size(MDP,1)
            MDP(m).o  = MDP(m).o(:,1:T);        % outcomes at 1,...,T
            MDP(m).s  = MDP(m).s(:,1:T);        % states   at 1,...,T
            MDP(m).u  = MDP(m).u(:,1:T);        % control  at 1,...,T - 1
        end
    end

end % end of loop over time


% loop over models to accumulate Dirichlet parameters and prepare outputs
%==========================================================================
for m = 1:size(MDP,1)

    % Smoothing of unchanging control states
    %----------------------------------------------------------------------
    for f = 1:Nf(m)
        if ~U{m}(f)
            for t = 1:T
                P{m,f,t} = P{m,f,T};
            end
        end
    end

    % Smoothing or backwards pass: replay
    %======================================================================
    if OPTIONS.B

        % reset parameters (concentration parameters): a
        %------------------------------------------------------------------
        for g = 1:Ng(m)
            if isfield(MDP,'a')
                a{m,g} = MDP(m).a{g};
            end
        end
        
        [Q,P,F]  = spm_backwards(O,P,Q,D,E,a,b,U,m);

        % update free energy
        %------------------------------------------------------------------
        MDP(m).F = F;

    end


    % prior Dirichlet parameters
    %======================================================================

    % over likelihood
    %----------------------------------------------------------------------
    for g = 1:Ng(m)
        if isfield(MDP,'a')
            pA{m,g} = MDP(m).a{g};
        end
        if isfield(MDP,'c')
            pC{m,g} = MDP(m).c{g};
        end
    end

    %  over hidden states
    %----------------------------------------------------------------------
    for f = 1:Nf(m)
        if isfield(MDP,'b')
            pB{m,f} = MDP(m).b{f};
        end
        if isfield(MDP,'d')
            pD{m,f} = MDP(m).d{f};
        end
        if isfield(MDP,'e')
            pE{m,f} = MDP(m).e{f};
        end
    end


    % learning - accumulate concentration parameters
    %======================================================================
    for t = 1:T

        % likelihood mapping from hidden states to outcomes: a
        %------------------------------------------------------------------
        if isfield(MDP(m),'a')
            for g = 1:Ng(m)
                da = spm_cross(O(m,g,t),Q{m,:,t});
                da = reshape(da,[No(m,g),Ns(m,:)]);
                da = da.*(A{m,g} > 0);
                MDP(m).a{g} = (MDP(m).a{g} + da)*eta;
            end
        end

        % mapping from hidden states to hidden states: b(u)
        %------------------------------------------------------------------
        if isfield(MDP,'b') && t < T
            for f = 1:Nf(m)
                db = spm_cross(spm_cross(Q{m,f,t + 1},Q{m,f,t}),P{m,f,t});
                db = db.*(MDP(m).b{f} > 0);
                MDP(m).b{f} = (MDP(m).b{f} + db)*eta;
            end
        end
    end

    % accumulation of prior preferences: (c)
    %----------------------------------------------------------------------
    if isfield(MDP,'c') && t < T
        for g = 1:Ng(m)
            dc = O{m,g,t + 1};
            dc = dc.*(MDP(m).c{g} > 0);
            MDP(m).c{g} = (MDP(m).c{g} + dc)*eta;
        end
    end

    % initial hidden states:
    %----------------------------------------------------------------------
    if isfield(MDP,'d')
        for f = 1:Nf(m)
            dd = Q{m,f,1};
            dd = dd.*(MDP(m).d{f} > 0);
            MDP(m).d{f} = (MDP(m).d{f} + dd)*eta;
        end
    end

    % initial hidden states:
    %----------------------------------------------------------------------
    if isfield(MDP,'e')
        for f = 1:Nf(m)
            de = P{m,f,1};
            de = de.*(MDP(m).e{f} > 0);
            MDP(m).e{f} = (MDP(m).e{f} + de)*eta;
        end
    end

    % (negative) free energy of parameters (complexity): outcome specific
    %======================================================================
    for g = 1:Ng(m)
        if isfield(MDP,'a')
            MDP(m).Fa(g) = - spm_KL_dir(MDP(m).a{g},pA{m,g});
        end
        if isfield(MDP,'c')
            MDP(m).Fc(f) = - spm_KL_dir(MDP(m).c{g},pC{g});
        end
    end

    % (negative) free energy of parameters: state specific
    %----------------------------------------------------------------------
    for f = 1:Nf(m)
        if isfield(MDP,'b')
            MDP(m).Fb(f) = - spm_KL_dir(MDP(m).b{f},pB{m,f});
        end
        if isfield(MDP,'d')
            MDP(m).Fd(f) = - spm_KL_dir(MDP(m).d{f},pD{m,f});
        end
        if isfield(MDP,'e')
            MDP(m).Fe(f) = - spm_KL_dir(MDP(m).e{f},pE{m,f});
        end
    end

    % reorganise posteriors for saving
    %======================================================================

    % states and paths
    %----------------------------------------------------------------------
    for t = 1:T
        for f = 1:Nf(m)
            X{m,f}(:,t) = Q{m,f,t};
            S{m,f}(:,t) = P{m,f,t};
        end
    end


    % simulated electrophysiological responses
    %======================================================================
    if OPTIONS.N

        % initialise simulated neuronal respsonses
        %--------------------------------------------------------------------------
        n     = 16;
        for f = 1:Nf(m)
            xn{m,f} = zeros(n,Ns(m,f),T,T);
        end

        % simulated dopamine (or cholinergic) responses: assuming a
        % monoexponential kernel
        %------------------------------------------------------------------
        h     = exp(-(0:(n - 1))/2);
        h     = h/sum(h);
        wn{m} = kron(w{m},ones(1,n));
        wn{m} = conv(wn{m},[spm_zeros(h) h],'same');
        dn{m} = gradient(wn{m}(:));


        % Belief updating about hidden states: assuming a kernel or impulse
        % response function with a cumulative gamma distribution
        %------------------------------------------------------------------
        h     = spm_Gcdf(0:(n - 1),n/4,1);
        for f = 1:Nf(m)
            for i = 1:Ns(m,f)
                for j = 1:T
                    for k = 1:T
                        if k == 1
                            h0 = 1/Ns(m,f);
                        else
                            h0 = sn{m,f}(i,j,k - 1);
                        end
                        ht     = sn{m,f}(i,j,k);
                        xn{m,f}(:,i,j,k) = h*(ht - h0) + h0;
                    end
                end
            end
        end

        % sum to one contraint
        %------------------------------------------------------------------
        for i = 1:n
            for j = 1:T
                for k = 1:T
                    xn{m,f}(i,:,j,k) = spm_norm(xn{m,f}(i,:,j,k)');
                end
            end
        end

        % belief updating about policies
        %------------------------------------------------------------------
        u0    = spm_softmax(ones(Np(m),1));
        for k = 1:Np(m)
            for t = 1:(T - 1)
                if t == 1
                    h0 = u0(k);
                else
                    h0 = R{m}(k,t - 1);
                end
                ht         = R{m}(k,t);
                j          = (1:n) + (t - 1)*n;
                un{m}(k,j) = (h*(ht - h0) + h0);
            end
        end

    end % if neural responses


    % assemble results and place in NDP structure
    %======================================================================
    MDP(m).T  = T;            % number of outcomes
    MDP(m).U  = V{m};         % policies
    MDP(m).R  = R{m};         % conditional expectations over policies
    MDP(m).C  = C(m,:);       % utility
    MDP(m).X  = X(m,:);       % conditional expectations over states
    MDP(m).P  = S(m,:);       % conditional expectations over controls
    MDP(m).O  = O(m,:,:);     % outcomes
    MDP(m).v  = v{m};         % expected free energy  over policies
    MDP(m).w  = w{m};         % precision of beliefs about policies

    MDP(m).O  = squeeze(MDP(m).O);
    MDP(m).P  = squeeze(MDP(m).P);

    if OPTIONS.N

        % neuronal responses
        %----------------------------------------------------------------------
        MDP(m).xn = xn(m,:);  % simulated neuronal encoding of states
        MDP(m).un = un{m};    % simulated neuronal encoding of policies
        MDP(m).wn = wn{m};    % simulated neuronal encoding of precision
        MDP(m).dn = dn{m};    % simulated dopamine responses (phasic)

    end

end % end loop over models (m)


% plot
%==========================================================================
if OPTIONS.P
    if ishandle(OPTIONS.P)
        figure(OPTIONS.P); clf
    else
        spm_figure('GetWin','MDP'); clf
    end
    spm_MDP_VB_trial(MDP(1))
end


% auxillary functions
%==========================================================================
function [G,P,F] = spm_forwards(O,P,A,B,C,H,W,t,T,N,m)
% deep tree search over policies or paths
%--------------------------------------------------------------------------
% FORMAT [G,Q,F] = spm_forwards(O,P,A,B,C,H,W,t,T,N,m)
% O{m,g,t} - cell array of outcome probabilities for modality g
% P{m,f,t}   cell array of priors over states
% A{m,g}   - likelihood mappings from hidden states
% B{m,f,k} - belief propagators (policy-dependent probability transitions)
% C{m,g}   - cost: log priors over future outcomes
% H{m,g}   - state dependent ambiguity
% W{m,g}   - state and outcome dependent novelty
% t        - current time point
% T        - time horizon
% N        - policy horizon
% m        - model or agent to update
%
% G(k)     - expected free energy over k policies
% Q{m,f,t} - posterior over states
% F        - variational free energy (negative or ELBO)
%
%  This subroutine performs a deep tree search over sequences of actions to
%  evaluate the expected free energy over policies or paths. Crucially, it
%  only searches likely policies under likely hidden states in the future.
%  This search is sophisticated; in the sense that posterior beliefs are
%  updated on the basis of future outcomes to evaluate the free energy
%  under each outcome. The resulting  average is then accumulated to
%  furnish a path integral of expected free energy for the next action.
%  This routine operates recursively by updating predictive posteriors over
%  hidden states and their most likely outcomes.
%__________________________________________________________________________


% Posterior over hidden states based on likelihood (A) and priors (P)
%==========================================================================
G        = zeros(size(B,3),1);                    % log priors over actions

% variational (Bayesian) belief updating and free energy (ELBO)
%--------------------------------------------------------------------------
[Q,F]    = spm_VBX(O(m,:,t),P(m,:,t),A(m,:),'full');
P(m,:,t) = Q;


% terminate search at time horizon
%--------------------------------------------------------------------------
if t == T || numel(G) == 1, return, end


% Expected free energy of subsequent action
%==========================================================================
for k = 1:size(B,3)                               % search over actions

    % (negative) expected free energy
    %----------------------------------------------------------------------
    for f = 1:size(B,2)
        Q{f,k} = B{m,f,k}*P{m,f,t};               % predictive posterior
    end

    for g  = 1:size(A,2)                          % for all modalities

        % predictive posterior and prior over outcomes
        %------------------------------------------------------------------
        qo   = spm_dot(A{m,g},Q(:,k));            % predictive outcomes
        po   = C{m,g};                            % predictive log prior

        % G(k): risk
        %-----------------------------------------------------------------
        G(k) = G(k) - qo'*(spm_log(qo) - po);

        % G(k): ambiguity
        %------------------------------------------------------------------
        if any(H{m,g},'all')
            G(k) = G(k) + spm_dot(H{m,g},Q(:,k));
        end

        % expected information gain (parameters: i.e., novelty)
        %------------------------------------------------------------------
        if any(W{m,g},'all')
            G(k) = G(k) + qo'*spm_dot(W{m,g},Q(:,k));
        end

    end
end

% deep (recursive) search over action sequences ( i.e., paths)
%==========================================================================
if t < N

    % probability over action (terminating search at a suitable threshold)
    %----------------------------------------------------------------------
    u     = spm_softmax(G);
    k     = u <= max(u)/16;
    u(k)  = 0;
    G(k)  = - 64;
    for k = 1:size(B,3)                      % search over actions
        q = spm_vec(spm_cross(Q(:,k)));      % vectorise posterior
        if u(k)                              % evaluating plausible paths

            %  evaluate expected free energy for plausible hidden states
            %--------------------------------------------------------------
            j     = find(q > max(q)/16);
            if length(j) > 16
                [q,j] = sort(q,'descend');
                j     = j(1:16);
            end
            for i = j(:)'

                % outcome probabilities under hidden state (i)
                %----------------------------------------------------------
                for g = 1:size(A,2)
                    O{m,g,t + 1} = A{m,g}(:,i);
                end

                % prior over subsequent action under this hidden state
                %----------------------------------------------------------
                P(m,:,t + 1) = Q(:,k);
                E      = spm_forwards(O,P,A,B,C,H,W,t + 1,T,N,m);

                % expected free energy marginalised over subsequent action
                %----------------------------------------------------------
                EFE(i) = spm_softmax(E)'*E;

            end

            % accumulate expected free energy marginalised over states
            %--------------------------------------------------------------
            G(k) = G(k) + EFE(j)*q(j);

        end % plausible paths

    end % search over actions


    % Predictive posterior over hidden states
    %----------------------------------------------------------------------
    u     = spm_softmax(G);
    for k = 1:size(B,2)
        for f = 1:size(B,2)
            R{f}  = u(k)*Q{f,k};
        end
    end
    P(m,:,t + 1) = R;

end


function [Q,P,F] = spm_backwards(O,P,Q,D,E,a,b,U,m)
% Backwards smoothing to evaluate posterior over initial states
%--------------------------------------------------------------------------
% O{m,g,t} - cell array of outcome probabilities for modality g
% P{m,k,t} - cell array of posteriors over paths
% Q{m,f,t} - cell array of posteriors over states
% D{m,f}   - cell array of priors over initial states
% E{m,k}   - belief propagators (action dependent probability transitions)
% a{m,g}   - likelihood tensor  (Dirichlet parameters)
% b{m,f}   - belief propagators (Dirichlet parameters)
% U{f}   - controllable factors
% m      - agent or model
%
% F      - Negative free energy (states, paths and parameters) ELBO
%
%  This subroutine performs Bayesian smoothing in the sense of a replay
%  using variational iterations to optimise posteriors over states, paths
%  and parameters, given the outcomes over an epoch. It effectively
%  implements the prior constraint that certain random variables (i.e., the
%  paths of uncontrollable factors and parameters) do not change with time
%__________________________________________________________________________

%  (iterative) variational scheme
%==========================================================================
T     = size(Q,3);

% for t = 1:T
%     for f = 1:numel(b(m,:))
%         Q{m,f,t} = spm_softmax(spm_log(Q{m,f,t})/512);
%     end
% end

% variational iterations
%--------------------------------------------------------------------------
tr    = @(A) pagetranspose(A);
Z     = -Inf;
for v = 1:4

    % initialise posterior Dirichlet parameters
    %----------------------------------------------------------------------
    F    = 0;                              % variational free energy (ELBO)
    qa   = a(m,:);
    qb   = b(m,:);

    % acccumulate posterior Dirichlet parameters
    %======================================================================
    for t = 1:T

        % likelihood mapping from hidden states to outcomes: a
        %------------------------------------------------------------------
        for g = 1:numel(qa)
            qa{g} = qa{g} + spm_cross(O{m,g,t},Q{m,:,t});
            qa{g} = qa{g}.*(a{g} > 0);
        end

        % mapping from hidden states to hidden states: b(u)
        %------------------------------------------------------------------
        if t < T
            for f = 1:numel(qb)
                qb{f} = qb{f} + spm_cross(spm_cross(Q{m,f,t + 1},Q{m,f,t}),P{m,f,t});
                qb{f} = qb{f}.*(b{f} > 0);
            end
        end

    end

%     % ELBO free energy of parameters (complexity) [NB: evaluated later]
%     %--------------------------------------------------------------------
%     for g = 1:numel(qa)
%         F = F - spm_KL_dir(qa{g},a{g});
%     end
%     for f = 1:numel(qb)
%         F = F - spm_KL_dir(qb{f},b{f});
%     end

    % inference (Bayesian filtering)
    %======================================================================
    for t = 1:T

        % accumulate likelihoods over modalities
        %------------------------------------------------------------------
        L     = 0;
        for g = 1:numel(qa)
            L = L + spm_dot(spm_psi(qa{g}),O{m,g,t});
        end

        % preclude numerical overflow of log likelihood
        %------------------------------------------------------------------
        L     = max(L,max(L,[],'all') - 8);
        
        % posterior over hidden states
        %------------------------------------------------------------------
        for f = 1:numel(qb)

            % log likelihood
            %--------------------------------------------------------------
            LL   = spm_vec(spm_dot(L,Q(m,:,t),f));

            % log prior: smoothing
            %--------------------------------------------------------------
            LP   = 0;
            if t == 1
                LP = LP + spm_log(D{m,f});
            end
            if t < T
                LP = LP + spm_dot(spm_psi(tr(qb{f})),P(m,f,t))*Q{m,f,t + 1};
            end
            if t > 1
                LP = LP + spm_dot(spm_psi(qb{f}),P(m,f,t - 1))*Q{m,f,t - 1};
            end

            % posterior
            %--------------------------------------------------------------
            Q{m,f,t} = spm_softmax(LL + LP);

            % ELBO free energy of states (accuracy and complexity)
            %--------------------------------------------------------------
            F     = F + Q{m,f,t}'*(LL + LP - spm_log(Q{m,f,t}));

        end

    end

    % beliefs about paths
    %======================================================================
    for f = 1:numel(qb)

        % beliefs about (changing) paths
        %------------------------------------------------------------------
        if U{m}(f)

            for t = 2:T

                % log likelihood of control states
                %----------------------------------------------------------
                LL = spm_dot(spm_dot(spm_psi(qb{f}),Q{m,f,t}),Q{m,f,t - 1});

                % prior over control state
                %----------------------------------------------------------
                LP = spm_log(P{m,f,t - 1});

                % posterior over control states
                %----------------------------------------------------------
                P{m,f,t - 1} = spm_softmax(LL + LP);

                % ELBO free energy of paths (complexity)
                %----------------------------------------------------------
                F  = F + P{m,f,t - 1}'*(LL + LP - spm_log(P{m,f,t - 1}));

            end

        else  % beliefs about (unchanging) paths

            % accumulate log likelihood
            %--------------------------------------------------------------
            LL    = 0;
            for t = 2:T
                LL = LL + spm_dot(spm_dot(spm_psi(qb{f}),Q{m,f,t}),Q{m,f,t - 1});
            end

            % prior over control state
            %--------------------------------------------------------------
            LP = spm_log(E{m,f});

            % posterior over control states
            %--------------------------------------------------------------
            PP    = spm_softmax(LL + LP);
            for t = 1:T
                P{m,f,t} = PP;
            end

            % ELBO free energy of paths (complexity)
            %--------------------------------------------------------------
            F  = F + PP'*(LL + LP - spm_log(PP));

        end
    end

    % convergence
    %======================================================================
    dF = F - Z;

    % % checks on ELBO
    %------------------------------------------------------------------------
%     if dF < 0
%         warning('ELBO decreasing in spm_smoothing')
%     end
%     if F > 0
%         warning('positive ELBO in spm_smoothing')
%     end
    if dF < 1/128
        break
    else
        Z = F;
    end

end

% NOTES:
%==========================================================================



return


function A  = spm_log(A)
% log of numeric array plus a small constant
%--------------------------------------------------------------------------
A           = log(A);
A(isinf(A)) = -32;


function A  = spm_norm(A)
% normalisation of a probability transition matrix (columns)
%--------------------------------------------------------------------------
A           = rdivide(A,sum(A,1));
A(isnan(A)) = 1/size(A,1);


function A  = spm_wnorm(A)
% summation of a probability transition matrix (columns)
% A = minus(log(A0),log(A)) + minus(1./A,1./A0) + minus(psi(A),psi(A0))
%   = minus(1./A,1./A0)/2 + ...
%--------------------------------------------------------------------------
A   = A + exp(-32);
A0  = sum(A,1);
A   = minus(log(A0),log(A)) + minus(1./A,1./A0) + minus(psi(A),psi(A0));
A   = 128*full(A);

function A  = spm_margin(A,i)
% marginalise a joint distribution
%--------------------------------------------------------------------------
if isvector(A), return, end
d     = 1:ndims(A);
d(i)  = [];
A     = sum(A,d(1));
for i = 2:numel(d)
    A = sum(A,d(i));
end
A     = squeeze(A);

return

function [M,MDP] = spm_MDP_get_M(MDP,T,Ng)
% FORMAT [M,MDP] = spm_MDP_get_M(MDP,T,Ng)
% returns an update matrix for multiple models
% MDP(m) - structure array of m MPDs
% T      - number of trials or updates
% Ng(m)  - number of output modalities for m-th MDP
%
% M      - update matrix for multiple models
% MDP(m) - structure array of m MPDs
%
% In some applications, the outcomes are generated by a particular model
% (to maximise free energy, based upon the posterior predictive density).
% The generating model is specified in the matrix MDP(m).n, with a row for
% each outcome modality, such that each row lists the index of the model
% responsible for generating outcomes.
%__________________________________________________________________________

for m = 1:size(MDP,1)

    % check size of outcome generating agent, as specified by MDP(m).n
    %----------------------------------------------------------------------
    if ~isfield(MDP(m),'n')
        MDP(m).n = zeros(Ng(m),T);
    elseif isempty(MDP(m).n)
        MDP(m).n = zeros(Ng(m),T);
    end
    if size(MDP(m).n,1) < Ng(m)
        MDP(m).n = repmat(MDP(m).n(1,:),Ng(m),1);
    end
    if size(MDP(m).n,2) < T
        MDP(m).n = repmat(MDP(m).n(:,1),1,T);
    end

    % mode of generating model (most frequent over outcome modalities)
    %----------------------------------------------------------------------
    n(m,:) = mode(MDP(m).n.*(MDP(m).n > 0),1);

end

% reorder list of model indices for each update
%--------------------------------------------------------------------------
n     = mode(n,1);
for t = 1:T
    if n(t) > 0
        M(t,:) = circshift((1:size(MDP,1)),[0 (1 - n(t))]);
    else
        M(t,:) = 1:size(MDP,1);
    end
end


return

function MDP = spm_MDP_update(MDP,OUT)
% FORMAT MDP = spm_MDP_update(MDP,OUT)
% moves Dirichlet parameters from OUT to MDP
% MDP - structure array (new)
% OUT - structure array (old)
%__________________________________________________________________________

% check for concentration parameters at this level
%--------------------------------------------------------------------------
try,  MDP.a = OUT.a; end
try,  MDP.b = OUT.b; end
try,  MDP.c = OUT.c; end
try,  MDP.d = OUT.d; end
try,  MDP.e = OUT.e; end

% check for concentration parameters at nested levels
%--------------------------------------------------------------------------
try,  MDP.MDP(1).a = OUT.mdp(end).a; end
try,  MDP.MDP(1).b = OUT.mdp(end).b; end
try,  MDP.MDP(1).c = OUT.mdp(end).c; end
try,  MDP.MDP(1).d = OUT.mdp(end).d; end
try,  MDP.MDP(1).e = OUT.mdp(end).e; end

return






%% NOTES: variational approximations to mapping from outcome space to
% approximate (variational) posterior using spm_VBX
%==========================================================================
Nf = [32,4,4];                            % number of levels per factor
Ng = [1024,2,8];                          % number of  outcomes per modality
T  = 128;                                 % temperature

% likelihood mapping
%--------------------------------------------------------------------------
for g = 1:numel(Ng)
    A{g} = rand([Ng(g),Nf]).^T;
    A{g} = bsxfun(@rdivide,A{g},sum(A{g},1));
end

% prior
%--------------------------------------------------------------------------
for f = 1:numel(Nf)
    P{f} = rand(Nf(f),1);
    P{f} = bsxfun(@rdivide,P{f},sum(P{f},1));
end

% outcomes
%--------------------------------------------------------------------------
for g = 1:numel(Ng)
    O{g} = rand(Ng(g),1);
    O{g} = bsxfun(@rdivide,O{g},sum(O{g},1));
end

% posterior and variational free energy from spm_VBX
%--------------------------------------------------------------------------
clf
[Q,F] = spm_VBX(O,P,A,'full'); disp(F)

subplot(2,2,1), bar(Q{1}), axis square, hold on
title({'first latent factor','iterative'})
subplot(2,2,2), bar(Q{2}), axis square, hold on
title({'second latent factor','iterative'})

[Q,F] = spm_VBX(O,P,A,'exact'); disp(F)

subplot(2,2,3), bar(Q{1}), axis square, hold off
title({'first latent factor','non-iterative'})
subplot(2,2,4), bar(Q{2}), axis square, hold off
title({'second latent factor','non-iterative'})

%% NOTES: on Dirichlet distributions: illustrating the impact of uncertainty
%==========================================================================
% about parameters
%--------------------------------------------------------------------------
a = (rand(8,1).^4)/64;

subplot(2,2,1)
plot(spm_dir_norm(a)), hold on, plot(spm_softmax(spm_psi(a))), hold off
title ('low Dirichlet counts')

a = a*256;

subplot(2,2,2)
plot(spm_dir_norm(a)), hold on, plot(spm_softmax(spm_psi(a))), hold off
title ('high Dirichlet counts')
legend({'expectation','variational'})

%% and the difference between the log of an expected Dirichlet
%==========================================================================
% distribution and the expected log
%--------------------------------------------------------------------------
K     = 4;
for i = 1:64
    a    = ones(K,1);
    a(1) = i/32;
    a0   = sum(a);

    ELOG(i) = psi(a(1)) - psi(a0);
    LOGE(i) = log(a(1)/a0);
    P(i)    = a(1);
end

plot(P,LOGE,P,ELOG)
legend({'log of expectation','expectation of log'})

%% ccomplexity costs and the precision of likelihoods
%==========================================================================
p     = 1/32;
a     = rand(120,1) > 3/4;
d     = rand(120,1) > 3/4;
a0    = a*0;
for i = 1:64
    A       = i*a + p;
    A0      = a0  + p;
    P       = A   + d;
    P0      = A0  + d;
    KL(i,1) = spm_KL_dir(P,A) + spm_KL_dir(A0,A0);
    KL(i,2) = spm_KL_dir(A,A) + spm_KL_dir(P0,A0);

end
subplot(2,1,1),plot(KL)
title({'rrelative complexity cost'})
legend({'assimilation','new state'})









