%% Load data
clear 
datadir = '/Users/byeolkim/dropbox/FAST/sync/data/public sharing';
savename = fullfile(datadir, 'figure2_b.mat');
load(savename);

%% TEST - fmri sess2, web sess 1, web sess 2
% predicted value : Markov features *  % weighted map of trained model
for i = 1:3
    yfit_test{i} = zscore(feature_test{i})*stats.other_output{1};      
    [r(i,1) p(i,1)] = corr(yfit_test{i}, y_actual{i});   % correlation btw the upper two
end

%% Figure 2B: correlation of training set 
saveornot = 0;
fontsize = 25;

close all;
h = plot_correlation(stats.Y, stats.yfit, 'robust', 'xlabel', '', 'ylabel', '');
set(gcf, 'position', [207 1116 237 229], 'color', 'w');
set(gca, 'fontsize', fontsize, 'xlim', [-1.8 3], 'ylim', [-1.8 2.8], 'linewidth', 1.8, 'tickdir', 'out', 'ticklength', [.02 .02]);
grid off;

[r, p] = corr(stats.Y, stats.yfit)
% 0.5086
for i = 1:(numel(h{1})-1)
    set(h{1}(i), 'MarkerFaceColor', [.4 .4 .4], 'Linewidth', 2, 'Color', 'w', 'MarkerSize', 10);
end
h{1}(end).Color = [33,102,172]./255;

%% Figure 2B: correlation of testing - MRI sess2
close all;
test_i = 1;
h = plot_correlation(y_actual{test_i}, yfit_test{test_i},'robust', 'xlabel', '', 'ylabel', '');
set(gcf, 'position', [207 1116 237 229], 'color', 'w');
set(gca, 'fontsize', fontsize, 'xlim', [-1.8 3], 'ylim', [-1.5 1.5], 'linewidth', 1.8, 'tickdir', 'out', 'ticklength', [.02 .02]);
grid off;

[r, p] = corr(y_actual{test_i}, yfit_test{test_i})
% 0.4810    0.0071
for i = 1:(numel(h{1})-1)
    set(h{1}(i), 'MarkerFaceColor', [.4 .4 .4], 'Linewidth', 2, 'Color', 'w', 'MarkerSize', 10);
end
h{1}(end).Color = [33,102,172]./255;

%% Figure 2B: correlation of testing - Web sess1 
close all;

test_i = 2;
h = plot_correlation(y_actual{test_i}, yfit_test{test_i},'robust', 'xlabel', '', 'ylabel', '');
set(gcf, 'position', [207 1116 237 229], 'color', 'w');
set(gca, 'fontsize', fontsize, 'xlim', [-1.9 4], 'ylim', [-2 1.7], 'linewidth', 1.8, 'tickdir', 'out', 'ticklength', [.02 .02]);
grid off;

[r, p] = corr(y_actual{test_i}, yfit_test{test_i})
% 0.3739
for i = 1:(numel(h{1})-1)
    set(h{1}(i), 'MarkerFaceColor', [.4 .4 .4], 'Linewidth', 2, 'Color', 'w', 'MarkerSize', 10);
end
h{1}(end).Color = [33,102,172]./255;

%% Figure 2B: correlation of testing - Web sess2 
close all;

test_i = 3;
h = plot_correlation(y_actual{test_i}, yfit_test{test_i},'robust', 'xlabel', '', 'ylabel', '');
set(gcf, 'position', [207 1116 237 229], 'color', 'w');
set(gca, 'fontsize', fontsize, 'xlim', [-5.5 11.5], 'ylim', [-1.5 1.5], 'linewidth', 1.8, 'tickdir', 'out', 'ticklength', [.02 .02]);
grid off;

[r, p] = corr(y_actual{test_i}, yfit_test{test_i})
% 0.5302
for i = 1:(numel(h{1})-1)
    set(h{1}(i), 'MarkerFaceColor', [.4 .4 .4], 'Linewidth', 2, 'Color', 'w', 'MarkerSize', 10);
end
h{1}(end).Color = [33,102,172]./255;

%% Figure 2C: show high ranking features 
clear selected_vars stability_pos stability_neg final_features final_stability final_weights non_zero_*;

% 1. select only stable features (>50% from cv)
for j = 1:size(stats.other_output_cv,1)
    selected_vars(:,j) = double(stats.other_output_cv{j,1}~=0).*sign(stats.other_output_cv{j,1});
end

stability_pos{i} = sum(selected_vars==1,2)./size(selected_vars,2);
stability_neg{i} = sum(selected_vars==-1,2)./size(selected_vars,2);
pos_neg = {'pos', 'neg'};

% 2. and sort the weights, get features that pos/neg in the final model
[s_val, non_zero_idx] = sort(stats.other_output{1});
non_zero_pos_idx_sorted{i}  = flipud(non_zero_idx(s_val>0));
non_zero_neg_idx_sorted{i}  = non_zero_idx(s_val<0);
interp_features_pos = non_zero_pos_idx_sorted{i};
interp_features_neg = non_zero_neg_idx_sorted{i};
fprintf(' >>>>  Features for model %s  >>>>', num2str(i))
for j = 1:numel(pos_neg)    % positive / negative weights
    eval(['final_features{i,j} = features_descrip_all(interp_features_' pos_neg{j} ');']);
    eval(['final_stability{i,j} = stability_' pos_neg{j} '{i}(interp_features_' pos_neg{j} ');']);
    eval(['final_weights{i,j} = stats.other_output{1}(interp_features_' pos_neg{j} ');']);
end
table(vertcat(final_features{i,:}), vertcat(final_weights{i,:}), vertcat(final_stability{i,:}), ...
    'VariableNames',{'selected features','weights','stability'})

%            selected features           weights     stability
%     _______________________________    ________    _________
%     {'variance_safe_threat'       }      0.4898     0.98387 
%     {'trans_prob_safe_threat_12'  }     0.28971     0.53226 
%     {'mean_time'                  }     0.15936           1 
%     {'mean_vividness'             }     0.15804     0.90323 
%     {'variance_valence'           }     0.13281     0.87097 
%     {'mean_self-related'          }    0.099148           1 
%     {'trans_prob_valence_31'      }    0.093331     0.98387 
%     {'steady_state_prob_valence_1'}    0.065401     0.45161 
%     {'trans_prob_valence_13'      }    -0.28922           1 
%     {'variance_self-related'      }    -0.24834     0.98387 
%     {'mean_valence'               }    -0.20123     0.98387 
%     {'trans_prob_valence_23'      }    -0.19935           1 
