% This script consists of all analyses.
% Frances Jin & Peter Zeidman, 2022

%% load data
load Data_for_Model.mat; % empirical data from 9 human patients
ns = length(D.d);

%% DCM model inversion with Variational Bayes for all subjects
% initialization
GCM = cell(ns,1);

% loop through all subjects (can change 'for' to 'parfor' for speed)
parfor s = 1:ns
    tic
    % Specify DCM
    DCM = J_dcm_model_specify(D.d{s})
    
    % Fit it to the data
    DCM = J_dcm_model_estimate(DCM,false);
    
    % Record output
    GCM{s} = DCM;    
    toc
end
save('GCM_nosology.mat','GCM');
%% Examining the multi-start for a specific subject
% run this section if you are interested in seeing the performance of the
% multi-start strategy.
s = 2; % specify a subject ID
DCM = J_dcm_model_specify(D.d{s});
DCM = J_dcm_model_estimate(DCM,false);

% examining the output of the Rayleigh parameter 
A2_output = [];
for iter = 1:length(DCM.multistartcheck)
    A2_output(iter,:) = [DCM.multistartcheck(iter).pE.A(2),DCM.multistartcheck(iter).Ep.A(2), DCM.multistartcheck(iter).F];
end
figure;set(gcf,'color','w');
subplot(1,2,1)
plot(A2_output(:,1),A2_output(:,2) , '.','Color','k');axis square; xlim([16 18.5]); ylim([16 18.5]);
subplot(1,2,2)
bar(A2_output(:,3));
%% Plot predicted and actual symptoms
load('GCM_nosology.mat');

figure;set(gcf,'color','w');
rows = 3;
cols = ns;

colours = {'#003153','#008C8C','#E85827','k'}; 
ylabels = {'Psychosis','Depression','Mania','Exogenous input'};

% Plot symptoms (one per row)
i = 1;
for symptom = 1:3
    for s = 1:ns % [1 5 9]
        subplot(rows,cols,i);
        
        % Plot data
        DCM = GCM{s};
        plot(DCM.y(:,symptom),'-.','Color',colours{symptom}, 'LineWidth',1.5);
        ylim([-0.25 1.25])
        hold on;
        
        % Plot prediction
        plot(DCM.Y(:,symptom),'Color',colours{symptom}, 'LineWidth',1.5);
        axis square; ylim([-0.25, 1.25]);
                
        % Decorate
        if symptom == 1
            title(['s' num2str(s)]);
        elseif symptom == 3
            xlabel('Time');
        end
        if s == 1
            ylabel(ylabels{symptom});
        end
        
        set(gca,'FontSize',12);
        
        i = i + 1;
        
    end    
end
%% Plot latent variables
close all
nv = 3; % Number of latent variables to plot

colours = {'#003153','#008C8C','#E85827','k'}; 
ylabels = {'Psychosis','Depression','Mania','Exogenous input'};
linestyles = {'-.','-','--'};

figure;set(gcf,'color','w');
rows = nv;
cols = 3; % or 'ns'

ax = [];

% Plot exogenous input
i = 1;
for v = 1:nv
    for sub = [1 5 9] % or '1:ns'
        subplot(rows,cols,i);
        DCM = GCM{sub};
        x = DCM.x;
        s = DCM.s;
        U = DCM.U;
        
        switch v
            case 1
                % Plot exogenous inputs
                plot(U,'Color',colours{end},'LineWidth',1.5);
                if sub == 1
                    ylabel(ylabels{end});
                end
                ylim([-0.3 0.4]);
                xlabel('Time');
                title(['s' num2str(sub)]);
            case 2
                % Pathophysiological states (level 2)
                ax(sub) = gca;
                plot3(x(:,1), x(:,2), x(:,3),'Color','k');
                hold on; plot3(x(1,1), x(1,2), x(1,3), '.','Color','b','MarkerSize',15);
                xlim([-20 18]);ylim([-30 10]);zlim([0 40]);
                %xlabel('x1');ylabel('x2');zlabel('x3');
                grid on               
            case 3
                % Psychopathology states (level 2)                
                for j = 1:size(DCM.s,2)
                    plot(s(:,j),'Color',colours{j},'LineStyle',linestyles{3}); hold on;
                end
                hold off
                if sub == 1
                    ylabel('Psychopathology'); 
                    legend(ylabels);
                end
                xlabel('Time');
               
        end
                
        set(gca,'FontSize',12);     
        
        axis square
        
        i = i + 1;    
    end
end

% Link rotation of 3D plots
Link = linkprop(ax,{'CameraUpVector', 'CameraPosition', 'CameraTarget'}); %, 'XLim', 'YLim', 'ZLim'
setappdata(gcf, 'StoreTheLink', Link);

%% PEB and plot the Bayesian Model Comparison

PEB = spm_dcm_peb(GCM);
BMA = spm_dcm_peb_bmc(PEB);

%% Plot estimated parameters
fields = {'A','B','x','t',};
rows = 5;
cols = 4;

% Get all parameters from all subjects (subjects x parameters)
Ep = cellfun(@(x) spm_vec(x.Ep)', GCM, 'UniformOutput', false);
Ep = cell2mat(Ep);

% Get all variances from all subjects (subjects x parameters)
Vp = cellfun(@(x) diag(x.Cp)', GCM, 'UniformOutput', false);
Vp = cell2mat(Vp);

spm_figure('GetWin','Parameters'); spm_clf;

p = 1;

for i = 1:length(fields)    
    % Identify indices and names of free parameters within this field
    pidx   = spm_find_pC(GCM{1},fields{i});
    pnames = spm_fieldindices(GCM{1}.Ep,pidx);
   
    % Plot
    for j = 1:length(pidx)
        subplot(rows,cols,p);
        spm_plot_ci(Ep(:,pidx(j)), diag(Vp(:,pidx(j))));
        if length(pidx) > 1
            title(pnames{j});
        else
            title(pnames);
        end
        axis square;
        p = p + 1;
    end 
end

%% Create simulated subjects and fit to data
rng(1);
groupsize = 12; % subjects per group

% Load empirical group causal models (GCM)
eGCM = load('GCM_nosology.mat');
eGCM = eGCM.GCM;

GCM = {};

% Generate models
groups = [];
for g = 1:ns
    % Empirical posteriors on which to base this group
    DCM = eGCM{g};
    Ep  = DCM.Ep;
    Cp  = DCM.Cp;
    M   = DCM.M;
    U   = DCM.U;
    
    for sub = 1:groupsize
        fprintf('Group %d subject %d\n',g,sub);
        
        % Sample parameters from empirical posterior
        P = spm_normrnd(spm_vec(Ep),full(Cp),1);
        P = spm_unvec(P,Ep);
        
        % Generate data
        y = spm_nosology_gen(P,M,U);
        
        % Specify new model
        DCM = J_dcm_model_specify(y);
        
        % Store generative parameters
        DCM.gP = P;
        
        % Store the DCM and which group this subject is in
        GCM{end+1,1} = DCM;
        groups(end+1,1) = g;
    end
end

% Fit to data
% loop through all subjects (can change 'for' to 'parfor' for speed)
parfor s = 1:length(GCM)   
    % Fit it to the data
    DCM = J_dcm_model_estimate(GCM{s},false);
    
    % Record output
    GCM{s} = DCM;    
end

save('GCM_simulated.mat','GCM','groups');
%% Plot face validity (parameters)
load('GCM_simulated.mat');
fields = {'A','B','x','t',};

% Get all estimated parameters from all subjects (subjects x parameters)
Ep = cellfun(@(x) spm_vec(x.Ep)', GCM, 'UniformOutput', false);
Ep = cell2mat(Ep);

% Get all generative parameters from all subjects (subjects x parameters)
gP = cellfun(@(x) spm_vec(x.gP)', GCM, 'UniformOutput', false);
gP = cell2mat(gP);

lims = [14 26
        -10 15
        -10 40
        122 132];
spm_figure('GetWin','Face validity (parameters)'); spm_clf;
rows = 3;
cols = 3;
for i = 1:length(fields) 
    % Identify indices and names of free parameters within this field
    pidx   = spm_find_pC(GCM{1},fields{i});
    
    % Get data
    x = gP(:,pidx);    
    y = Ep(:,pidx);
    x = x(:);
    y = y(:);
    
    % Compute correlation
    R = corrcoef(x,y);
    R = R(1,2);
    
    % Plot diagonal line
    subplot(rows,cols,i);
    line([lims(i,1) lims(i,2)],[lims(i,1) lims(i,2)],'Color','k','LineStyle','--');
    hold on
    
    % Plot data
    plot(x, y, '.','Color','k');
    title(sprintf('%s\nr = %2.2f',fields{i},R));        
    
    % Line of best fit 
    h=lsline;  
    
    % Decorate plot
    xlim(lims(i,:));
    ylim(lims(i,:));
    axis square          
    set(h,'Color','r');
    hold off;    
    box on
    xlabel('Generative'); ylabel('Estimated');
    set(gca,'FontSize',12);
end
%% Compute face validity (second level model comparison, slow)
load('GCM_simulated.mat');

% Get first level prior expectations to use as the second level prior 
% expectations for PEB (we need to do this because all subjects will have
% different priors to each another, due to the multi-start algorithm)
dcm = J_dcm_model_specify(GCM{1}.y);
pE  = dcm.M.pE;

groupsize = sum(groups==1); % subjects per group
ngroups   = max(groups);

X = kron(eye(ngroups),ones(groupsize,1));
X = X(:,2:end);
X = [ones(length(GCM),1) X];

F = [];
parfor s = 1:size(GCM,1)
    for g = 1:ngroups
        % Construct design matrix
        Xb = X;
        Xb(s,2:end) = 0;        
        if g > 1
            Xb(s,g) = 1;
        end
        
        M = struct();
        M.X  = Xb;
        M.bE = pE;
        M.Q  = 'fields';
        M.maxit = 128;
        
        PEB = spm_dcm_peb(GCM,M,{'A','B','C','x','t'});
        F(s,g) = PEB.F;
    end
end

P = F-min(F,[],2);
P = spm_softmax(P')';

save('face_validity.mat','F','P','X');
%% Plot face validity (second level model comparison)
close all
load('face_validity.mat');

spm_figure('GetWin','Face validity (models)');
spm_clf;

subplot(2,2,1);
imagesc(X); colormap parula;colorbar
ylabel('Subject'); xlabel('Group');
set(gca,'FontSize',12);
axis square;
title('Design matrix','FontSize',14);

subplot(2,2,2);
imagesc(P); 
ylabel('Subject'); xlabel('Group');
colormap parula;colorbar
set(gca,'FontSize',12);
axis square;
title('Confusion matrix','FontSize',14);
%% Compute LOO analysis - can we distinguish group 1 from group 5? (slow)
% the two groups were selected arbitrarily
load('GCM_simulated.mat');

% Get first level prior expectations to use as the second level prior 
% expectations for PEB (we need to do this because all subjects will have
% different priors to each another, due to the multi-start algorithm)
dcm = J_dcm_model_specify(GCM{1}.y);
pE  = dcm.M.pE;

% Limit to subjects in group 1 and 5
idx = find(groups == 1 | groups == 5);
GCM = GCM(idx);
groups = groups(idx);

% Build design matrix
X = [ones(length(GCM),1) double(groups == 5)];

% PEB options
M = struct();
M.X  = X;
M.bE = pE;
M.Q  = 'fields';
M.maxit = 128;

[qE,qC,Q] = spm_dcm_loo(GCM,M,{'A'});

save('LOO_nosology_2groups.mat','qE','qC','Q','M','idx');

%% Plot LOO analysis results
load('LOO_nosology_2groups.mat');

% classical inference on classification accuracy
[T,df] = spm_ancova(M.X(:,1:2),[],qE(:),[0;1]);
r      = corrcoef(qE(:),M.X(:,2));
r      = full(r(1,2));
if isnan(T)
    p = NaN;
else
    p = 1 - spm_Tcdf(T,df(2));
end
str = sprintf('corr(df:%-2.0f) = %-0.2f: p = %-0.5f',df(2),r,p);

spm_figure('GetWin','LOO'); spm_clf;
subplot(2,2,1);
plot(M.X(:,2),qE,'.','Markersize',10,'Color','k')
xlabel('group effect'), ylabel('estimate')
set(gca,'FontSize',12);
title(str,'FontSize',14)
axis square
xlim([-0.1 1.1]);

%% LOO analysis for all 36 pairs (9 group)
load('GCM_simulated.mat');

ngroup = 9;
for ii = 1:ngroup
    for iii = ii+1:ngroup
        if iii ~= ii
            
            idx = find(groups == ii | groups == iii);
            gGCM = GCM(idx);
            ggroups = groups(idx);
            
            % Build design matrix
            X = [ones(length(gGCM),1) double(ggroups == iii)];
            
            % PEB options
            M = struct();
            M.X  = X;
            M.bE = pE;
            M.Q  = 'fields';
            M.maxit = 128;
            
            [qE,qC,Q] = spm_dcm_loo(gGCM,M,{'A'});
            LOO_all_pairs.qE{ii,iii} = qE;
            LOO_all_pairs.qC{ii,iii} = qC;
            LOO_all_pairs.Q{ii,iii} = Q;
            LOO_all_pairs.M{ii,iii} = M;
            LOO_all_pairs.idx{ii,iii} = idx;
            clear idx gGCM X M qE qC Q
 
        end
    end
end
save('LOO_nosology_allpairs.mat','LOO_all_pairs');

%% Plot LOO analysis results for all pairs
load('LOO_nosology_allpairs.mat');
% classical inference on classification accuracy
ngroup = 9;
corr_matrix = zeros(9);
for ii = 1:ngroup
    for iii = ii+1:ngroup
        [T,df] = spm_ancova(LOO_all_pairs.M{ii,iii}.X(:,1:2),[],LOO_all_pairs.qE{ii,iii}(:),[0;1]);
        r      = corrcoef(LOO_all_pairs.qE{ii,iii}(:),LOO_all_pairs.M{ii,iii}.X(:,2));
        r      = full(r(1,2));
        corr_matrix(ii,iii) = r;
    end
end
figure;set(gcf,'color','w');
subplot(2,2,1);
imagesc(corr_matrix); 
[R, C] = ndgrid(1:size(corr_matrix,1), 1:size(corr_matrix,2));
R = R(:); C = C(:) - 1/4;
vals = round(corr_matrix(:),1);
mask = vals <= 0.4;
text(C(mask), R(mask), string(vals(mask)), 'color', 'w')
text(C(~mask), R(~mask), string(vals(~mask)), 'color', 'k')
ylabel('Group'); xlabel('Group');
colormap parula;colorbar
set(gca,'FontSize',12);