function Q = spm_MDP_VB_game_SAD(MDP)
% auxiliary plotting routine for spm_MDP_VB - multiple trials
% FORMAT Q = spm_MDP_VB_game(MDP)
%
% MDP.P(M,T)      - probability of emitting action 1,...,M at time 1,...,T
% MDP.X           - conditional expectations over hidden states
% MDP.R           - conditional expectations over policies
% MDP.O(O,T)      - a sparse matrix encoding outcomes at time 1,...,T
% MDP.S(N,T)      - a sparse matrix encoding states at time 1,...,T
% MDP.U(M,T)      - a sparse matrix encoding action at time 1,...,T
% MDP.W(1,T)      - posterior expectations of precision
%
% MDP.xn  = Xn    - simulated neuronal encoding of policies
% MDP.wn  = wn    - simulated neuronal encoding of precision
% MDP.da  = dn    - simulated dopamine responses (deconvolved)
% MDP.rt  = rt    - simulated dopamine responses (deconvolved)
%
% returns summary of performance:
%
%     Q.X  = x    - expected hidden states
%     Q.R  = u    - final policy expectations
%     Q.S  = s    - initial hidden states
%     Q.O  = o    - final outcomes
%     Q.p  = p    - performance
%     Q.q  = q    - reaction times
%
% please see spm_MDP_VB
%__________________________________________________________________________
% Copyright (C) 2005 Wellcome Trust Centre for Neuroimaging

% Karl Friston
% $Id: spm_MDP_VB_game.m 7766 2020-01-05 21:37:39Z karl $

% numbers of transitions, policies and states
%--------------------------------------------------------------------------
if iscell(MDP(1).X)
    Nf = numel(MDP(1).B);                 % number of hidden state factors
    Ng = numel(MDP(1).A);                 % number of outcome factors
else
    Nf = 1;
    Ng = 1;
end


% graphics
%==========================================================================
Nt    = numel(MDP);                      % number of trials
Ne    = MDP(1).T;                        % number of outcomes per trial
Np    = size(MDP(1).V,2);                % number of policies
for i = 1:Nt
    
    % assemble expectations of hidden states and outcomes
    %----------------------------------------------------------------------
    for j = 1:Ne
        for k = 1:Ne
            for f = 1:Nf
                try
                    x{f}{i,1}{k,j} = gradient(MDP(i).xn{f}(:,:,j,k)')';
                catch
                    x{f}{i,1}{k,j} = gradient(MDP(i).xn(:,:,j,k)')';
                end
            end
        end
    end
    s(:,i) = MDP(i).s(:,1);
    o(:,i) = MDP(i).o(4,:)';
    v(1:2,i) = MDP(i).P(1:2,1)';
    v(3:4,i) = MDP(i).P(3:4,2)';
    v(5:6,i) = MDP(i).P(5:6,3)';
    act1(:,i) = MDP(i).u(1,1);
    act2(:,i) = MDP(i).u(1,2);
    act3(:,i) = MDP(i).u(1,3);
    w(:,i) = mean(MDP(i).dn,2);
    
    % assemble context learning
    %----------------------------------------------------------------------
    for f = 1:Nf
        try
            try
                D = MDP(i).d{f};
            catch
                D = MDP(i).D{f};
            end
        catch
            try
                D = MDP(i).d;
            catch
                D = MDP(i).D;
            end
        end
        d{f}(:,i) = D/sum(D);
    end
    
    % assemble performance
    %----------------------------------------------------------------------
    p(i)  = 0;
    for g = 1:Ng
        try
            U = spm_softmax(MDP(i).C{g});
        catch
            U = spm_softmax(MDP(i).C);
        end
        for t = 1:Ne
            p(i) = p(i) + log(U(MDP(i).o(g,t),t))/Ne;
        end
    end
end

% assemble output structure if required
%--------------------------------------------------------------------------
if nargout
    Q.X  = x;            % expected hidden states
    Q.R  = v;            % final policy expectations
    Q.S  = s;            % inital hidden states
    Q.O  = o;            % final outcomes
    Q.p  = p;            % performance
    Q.q  = q;            % reaction times

    return
end


% Initial states and expected policies (habit in red)
%--------------------------------------------------------------------------
col   = {'r.','g.','b.','c.','m.','k.'};
t     = 1:Nt;

subplot(3,1,1)
if Nt < 64
    MarkerSize = 24;
else
    MarkerSize = 8;
end

image(64*(1 - v)),  hold on
colorbar off
plot(act1,col{4},'MarkerSize',MarkerSize)
plot(act2,col{4},'MarkerSize',MarkerSize)
plot(act3,col{4},'MarkerSize',MarkerSize)

% try
%     plot(Np*(1 - v(Np,:)),'r')
% end
% try
%     E = spm_softmax(spm_cat({MDP.e}));
%     plot(Np*(1 - E(end,:)),'r:')
% end

title('Action selection and action probabilities')
xlabel('Trial'),ylabel('Action'), hold off
yticklabels({'Attend to Env','Attend to Self','Stay','Escape','Reflection','Rumination'})

% Performance
%--------------------------------------------------------------------------
subplot(3,1,2), bar(p,'k'),   hold on

for i = 1:size(o,2) 
      j(i,1) = MDP(i).o(3,3);

 if j(i,1) == 1
        jj(i,1) = 1; % mishap
 else
        jj(i,1) = 0;
 end
 if j(i,1) == 2
        nut(i,1) = 2; % nutural
 else
        nut(i,1) = 0;
 end
end

blackdots(i)=0;

plot((j), col{3}, 'MarkerSize',MarkerSize);
plot((jj), col{1}, 'MarkerSize',MarkerSize);
plot((nut), col{6}, 'MarkerSize',MarkerSize);
plot((blackdots),col{6}, 'MarkerSize',MarkerSize);

title('Perceived Social Consequences and Free energies')
ylabel('Value and Perceived Social Consequences'), spm_axis tight, hold off, box off
set(gca,'YTick',[-4:1:3])
yticklabels({'','','','Free Energy','','Mishap','Neutral','Positive'})

% learning - D
%--------------------------------------------------------------------------
subplot(3,1,3)

for f = 2
    subplot(5*Nf,1,Nf*4 + f), image(64*(1 - d{f}))
    title('Context Learning')
    x_ticks = [1, 20, 40, 60, 80, 100];
    set(gca,'XTick',x_ticks);
    yticklabels({'Safe','Dangerous'})
    
end


