%% Supplementary MATLAB code for:
% Towards quantitative cognitive-behavioral models: An active inference account of social anxiety disorder

clear, close all
rng('shuffle')

%% Simulation Setting

PriorSafe = 0.6; % 0-1, higher = stronger explicit prior belief to be in a safe environment
SEff = 0.8;      % 0-1, higher = higher confidence to be respected in social encounters
PostRum = 0;     % 0-infinity, higher = higher preference for rumiantion
SAttn = 0;       % 0-infinity, higher = higher preference for heightened self-focused attention

Learning = 1; % 0 = single trial; 1 = simulate learning
N = 100;      % number of social encounters

%% Model Specification

%--------------------------------------------------------------------------
% Initial State Priors (D)
%--------------------------------------------------------------------------

% action states: 
D{1} = [1 0 0 0 0 0 0]'; % 'start' 'attend to environment' 'attend to self' % 'stay' 'escape' 'realistic reflection' 'rumination'

% social context: 
D{2} = [0 1]';  % 'safe' 'threatening'

%--------------------------------------------------------------------------
% Agent's Explicit Beliefs ('d')
%--------------------------------------------------------------------------

d = D;

d{2} = [PriorSafe 1-PriorSafe]'; % 'safe' 'threatening'

d{1} = d{1}*50;
d{2} = d{2}*5; 

% Learning in active inference is formulated using a Dirichlet-categorical model, 
% with the distribution in d & e vectors being Dirichlet distribution. 

% In simple terms, the belief updating process involves adding counts to the 
% parameters encodinhg the relevant state/policy after obtaining the observations.
% Here, higher values of the parameters in d & e vectors
% indicate higher confidence ending up in the corresponding state (or policy).
% These distributions will be passed through a softmax function during within 
% trial inference, such that they remain the same shape while still add up to 1.

%--------------------------------------------------------------------------
% State-Observation Mappings (A Matrices)
%--------------------------------------------------------------------------

Nf    = numel(D);
for f = 1:Nf
    Ns(f) = numel(D{f});
end

No    = [2 3 3 7]; 
% possible outcomes per modality: 
% {interoception, exteroception, perceived social consequences, actions executed}

Ng    = numel(No);
for g = 1:Ng
    A{g} = zeros([No(g),Ns]);
end

%-----------------------------------------------------------------------
% Interoception (arousal level) {'high' 'low'}

% social context: safe
% 'start' 'attend to environment' 'attend to self' 'stay' 'escape' 'realistic reflection' 'rumination'
A{1}(:,:,1) =   [0   0.2 0.7   0.2 0.2   0.2 0.9;  % high
                 1   0.8 0.3   0.8 0.8   0.8 0.1]; % low
               
% social context: threatening          
A{1}(:,:,2) =   [0   0.6 0.8   0.6 0.2   0.6 0.9;  % high
                 1   0.4 0.2   0.4 0.8   0.4 0.1]; % low
%-----------------------------------------------------------------------
% Exteroception {'judgemental' 'friendly' 'neutral'}
             
% social context: safe 
A{2}(:,:,1) =   [0   0.04 0    0.2 0    0.04 0.5;  % judgemental
                 0   0.96 0    0.8 0    0.96 0.5;  % friendly
                 1   0    1      0 1      0    0]; % neutral

% social context: threatening       
A{2}(:,:,2) =   [0   0.96 0    0.8 0    0.96  0.96;  % judgemental
                 0   0.04 0    0.2 0    0.04  0.04;  % friendly
                 1      0 1      0 1       0     0]; % neutral

% This mapping is highly precise compared to A{1}: it really informs social context

%-----------------------------------------------------------------------
% Perceived social consequences {'mocked' 'neutral' 'respected'}

% social context: safe 
A{3}(:,:,1) =   [0   0 0   1-SEff 0   1-SEff 0.8;   % mocked
                 1   1 1   0      1   1      0;     % neutral
                 0   0 0   SEff   0   SEff   0.2];  % respected

% social context: threatening
A{3}(:,:,2) =   [0   0 0   1-SEff/5 0  1-SEff/5 0.96;  % mocked
                 1   1 1   0        1  1        0;     % neutral
                 0   0 0   SEff/5   0  SEff/5   0.04]; % respected

%-----------------------------------------------------------------------
% Observed actions

for i = 1:2

    A{4}(:,:,i) = eye(7);
    
end

%--------------------------------------------------------------------------
%Action-Dependent State Transitions (B Matrices)
%--------------------------------------------------------------------------
for f = 1:Nf
    B{f} = eye(Ns(f));
end
 
    B{1}(:,:,1) = [0 0 0 0 0 0 0
                   1 1 1 1 1 1 1
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0]; % observe: Attend to environment
               
    B{1}(:,:,2) = [0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   1 1 1 1 1 1 1
                   0 0 0 0 0 0 0 
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0]; % observe: Attend to self
   
    B{1}(:,:,3) = [0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   1 1 1 1 1 1 1
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0]; % choose to stay

    B{1}(:,:,4) = [0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   1 1 1 1 1 1 1
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0]; % choose to escape

    B{1}(:,:,5) = [0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0 
                   0 0 0 0 0 0 0
                   1 1 1 1 1 1 1
                   0 0 0 0 0 0 0]; % realistic reflection

    B{1}(:,:,6) = [0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0 
                   0 0 0 0 0 0 0
                   0 0 0 0 0 0 0
                   1 1 1 1 1 1 1]; % rumination
    
%--------------------------------------------------------------------------
%Allowable Policies (Pi)
%--------------------------------------------------------------------------
T = 4; % number of timesteps

Np        = 8; % number of policies

V         = ones(T-1,Np,Nf);

V(:,:,1) = [1 3 5;     % 1. attn-env - stay   - ref
            1 3 6;     % 2. attn-env - stay   - rum
            
            1 4 5;     % 3. attn-env - escape - ref
            1 4 6;     % 4. attn-env - escape - rum

            2 3 5;     % 5. attn-self - stay   - ref
            2 3 6;     % 6. attn-self - stay   - rum

            2 4 5;     % 7. attn-self - escape - ref
            2 4 6]';   % 8. attn-self - escape - rum
        

%Prior over policies (E)
    
E = [1 1 1 1 1 1 1 1]';

e = E;
e = e*512; % Dirichlet distribution, see annotation for d.

%--------------------------------------------------------------------------
%Preferences (C vectors)
%--------------------------------------------------------------------------
C{1}      = zeros(No(1),T);
C{2}      = zeros(No(2),T);
C{3}      = zeros(No(3),T);
C{4}      = zeros(No(4),T);

C{1}(1,:) = -0.1;    % dysprefer high arousal
C{2}(1,:) = -0.2;    % dysprefer judgemental cues
C{3}(1,:) = -2;    % strongly dysprefer to be mocked
C{3}(3,:) =  0.9;    % prefer to be respected

C{4}(3,:) = SAttn;   % preference for self-focused attention 
C{4}(7,:) = PostRum; % preference for rumination

%% Define MDP Structure

mdp.T = T;                      % number of time points
mdp.V = V;                      % allowable policies
mdp.A = A;                      % state-outcome mappings
mdp.B = B;                      % transition probabilities
mdp.C = C;                      % preferred outcomes
mdp.D = D;                      % prior over initial states
mdp.E = E;                      % prior over policies

mdp.d = d;                      % agent's belief to be learned
mdp.e = e;                      % agent's belief to be learned

mdp.beta    = 8;                % policy precision
mdp.alpha   = 32;               % action precision (inverse temperature)
mdp.eta     = 1;                % learning rate 
mdp.erp     = 1;                % degree of belief resetting at each timestep
mdp.tau     = 12;               % update time constant

label.factor{1}   = 'Action States';   label.name{1}    = {'start' 'attn-env' 'attn-self' 'stay' 'escape' 'ref' 'rum'};
label.factor{2}   = 'True Social Context';     label.name{2}    = {'safe','threating'};

label.modality{1} = 'Interocept';    label.outcome{1} = {'high','low'};
label.modality{2} = 'Exterocept';    label.outcome{2} = {'judgemental','friendly','neutral'};
label.modality{3} = 'Perceived Social Consequence';    label.outcome{3} = {'mocked','neutral','respected'};
label.modality{4} = 'Actions';       label.outcome{4} = {'start','attn-env','attn-self','stay','esc','ref','rum'}; 

label.action{1} = {'attn-env','attn-self','stay','esc','ref','rum'};

mdp.label = label;

MDP = spm_MDP_check(mdp);

MDP = spm_MDP_VB_X(MDP);

% Illustrate single trial before learning

spm_figure('GetWin','Single Trial Behaviour'); clf
spm_MDP_VB_trial(MDP,1:2,1:4);


%% Simulate Learning

if Learning == 1

[MDP2(1:N)] = deal(mdp);

    for i = 1:30
        MDP2(i).D{2}   = [0.1 0.9]'; % Start in the 'threatening' context 
        % mdp.eta = 1; % uncross for simulation 6: higher learning rate in
%         threatening context
    end
% 
    for i = 31:N
        MDP2(i).D{2}   = [0.9 0.1]'; % Switch to 'safe' context
        % mdp.eta = 0.2; % uncross for simulation 6: lower learning rate in
%         safe context
    end

MDP2 = spm_MDP_VB_X(MDP2);
% changes in explicit belief about the environment

% plotting changes in agent's d & e matrices
%--------------------------------------------------------------------------
for i = 1:N
    d_evo(1:2,i) = MDP2(i).d{2}(:,1);
    % e_evo(1:8,i) = MDP2(i).e(:,1);
end

dsm = spm_softmax(d_evo,.1);
for i = 1:N
    x = linspace(1,N,N);
    y1(1,i)= dsm(1,i,1);
    y2(1,i)= dsm(2,i,1);
end

figure; hold on
ylim([0,1])
plot(x,y1,'LineWidth',3,'Color','#0072BD');
line([1,30],[0.1,0.1],'LineWidth',4,'Color','#000000','LineStyle','--');
line([31,100],[0.9,0.9],'LineWidth',4,'Color','#000000','LineStyle','--');
% text(31,0.93,'true p(context=safe)','Color','#000000');
% text(2,0.13,'true p(context=safe)','Color','#000000');
xlabel('Number of Trials');
legend({'p=safe(inferred)'});
title('Explicit beliefs over time in exposure');

%--------------------------------------------------------------------------
% %% Sample behaviour under beliefs at a given time point during exposure
    Test_N = 90;  % time point to sample (0-100)
    mdp.D = MDP2(Test_N).D;
    mdp.d = MDP2(Test_N).d;
    mdp.e = MDP2(Test_N).e;

    MDP3 = spm_MDP_check(mdp);
    MDP3 = spm_MDP_VB_X(MDP3);

    spm_figure('GetWin','post-exposure behaviour'); clf
    spm_MDP_VB_trial(MDP3,1:2,1:4);
    spm_figure('GetWin','learning'); clf
    spm_MDP_VB_game_SAD(MDP2);
end
