% Create an overview plot featuring the results of the multivariate PLS
% comparing spectral changes during the stimulus period under load

clear all; cla; clc;

currentFile = mfilename('fullpath');
[pathstr,~,~] = fileparts(currentFile);
cd(fullfile(pathstr,'..'))
rootpath = pwd;

pn.data         = fullfile(rootpath, 'data', 'stats');
pn.figures      = fullfile(rootpath, 'figures');
pn.tools        = fullfile(rootpath, 'tools');
    addpath(genpath(fullfile(pn.tools, '[MEG]PLS', 'MEGPLS_PIPELINE_v2.02b')))
    addpath(fullfile(pn.tools, 'fieldtrip')); ft_defaults;
    addpath(genpath(fullfile(pn.tools, 'RainCloudPlots')));
    addpath(fullfile(pn.tools, 'BrewerMap'));
    addpath(fullfile(pn.tools, 'winsor'));

% set custom colormap
cBrew = brewermap(500,'RdBu');
cBrew = flipud(cBrew);
colormap(cBrew)

load(fullfile(pn.data, 'd2_taskpls_erf.mat'),...
    'stat', 'result', 'lvdat', 'lv_evt_list', 'num_chans', 'num_freqs', 'num_time')
load(fullfile(rootpath, 'data','erp', ['sub-001_erp.mat']));
elec = erp.scene_category{1}.elec;

result.perm_result.sprob

indLV = 1;

lvdat = reshape(result.boot_result.compare_u(:,indLV), num_chans, num_freqs, num_time);
stat.prob = lvdat;
stat.mask = lvdat > 3 | lvdat < -3;

% maskNaN = double(stat.mask);
% maskNaN(maskNaN==0) = NaN;

%% invert solution

stat.mask = stat.mask;
stat.prob = stat.prob.*-1;
result.vsc = result.vsc.*-1;
result.usc = result.usc.*-1;

%% plot BSR loadings

h = figure('units','centimeter','position',[0 0 15 10]);
set(gcf,'renderer','Painters')
statsPlot = [];
%statsPlot = cat(1, statsPlot,squeeze(nanmean(stat.prob(1:64,:,:),1)));
statsPlot = cat(1, statsPlot,squeeze(nanmean(stat.mask(1:64,:,:).*stat.prob(1:64,:,:),1)));
range_plot = [-1 3.5];
imagesc(stat.time,[],statsPlot, range_plot);
set(gca,'Ydir','Normal');
% set custom colormap
ncol = 2000; cBrew = brewermap(ncol,'RdBu'); cBrew = flipud(cBrew);
% adjust for unequal range
reldev = min(abs(range_plot))./max(abs(range_plot));
if abs(range_plot(2))<max(abs(range_plot))
    nadjust = [1:ncol/2, ncol/2+round(linspace(1, ncol/2, reldev*(ncol/2)))];
elseif abs(range_plot(1))<max(abs(range_plot))
    nadjust = [round(linspace(1, ncol/2, reldev*(ncol/2))),ncol/2:ncol];
else
    nadjust = 1:ncol;
end
cBrew = cBrew(nadjust,:);
colormap(cBrew); colorbar;
xlabel('Time [s from stim onset]'); ylabel('Frequency (Hz)');
title({'ERF changes'; ['p = ', num2str(round(result.perm_result.sprob(indLV),4))]})
set(findall(gcf,'-property','FontSize'),'FontSize',18)
set(gca, 'YTickLabels', round(stat.freq(get(gca, 'YTick')),0));

figureName = ['d2_erf_bsr'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');

%% plot multivariate topographies

% set custom colormap
cBrew = brewermap(500,'RdBu');
cBrew = flipud(cBrew);

h = figure('units','centimeters','position',[0 0 10 10]);
set(gcf,'renderer','Painters')

cfg = [];
cfg.layout = 'biosemi64.lay';
cfg.parameter = 'powspctrm';
cfg.comment = 'no';
cfg.colormap = cBrew;
cfg.colorbar = 'EastOutside';

plotData = [];
plotData.label = elec.label; % {1 x N}
plotData.dimord = 'chan';
cfg.zlim = [-4 4]; 
cfg.figure = h;
plotData.powspctrm = squeeze(nanmean(nanmean(stat.mask(:,stat.freq>8 & stat.freq<20,stat.time>1 & stat.time<2).*...
    stat.prob(:,stat.freq>8 & stat.freq<20,stat.time>1 & stat.time<2),3),2)); 
[~, sortind] = sort(plotData.powspctrm, 'descend');
cfg.marker = 'off'; 
cfg.highlight = 'yes';
cfg.highlightchannel = plotData.label(sortind(1:6)); %26, 20, 30, 31, 21, 63
cfg.highlightcolor = [0 0 0];
cfg.highlightsymbol = '.';
cfg.highlightsize = 30;
ft_topoplotER(cfg,plotData);
cb = colorbar('location', 'EastOutside'); set(get(cb,'ylabel'),'string','Mean BSR');
title('AlphaBeta 1-2s')

figureName = ['d2_erf_topo1'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');

%% 'Theta 200 - 500 ms'

h = figure('units','centimeters','position',[0 0 10 10]);
set(gcf,'renderer','Painters')

cfg = [];
cfg.layout = 'biosemi64.lay';
cfg.parameter = 'powspctrm';
cfg.comment = 'no';
cfg.colormap = cBrew;
cfg.colorbar = 'EastOutside';

plotData = [];
plotData.label = elec.label; % {1 x N}
plotData.dimord = 'chan';
cfg.zlim = [-1.5 1.5]; 
cfg.figure = h;
plotData.powspctrm = squeeze(nanmean(nanmean(stat.mask(:,stat.freq>2 &stat.freq<7,stat.time>.3 & stat.time<.5).*...
    stat.prob(:,stat.freq>2 & stat.freq<7,stat.time>.3 & stat.time<.5),3),2)); 
[~, sortind] = sort(plotData.powspctrm, 'descend');
cfg.marker = 'off'; 
cfg.highlight = 'yes';
cfg.highlightchannel = plotData.label(sortind(1:2)); % 38, 47
cfg.highlightcolor = [0 0 0];
cfg.highlightsymbol = '.';
cfg.highlightsize = 30;
ft_topoplotER(cfg,plotData);
cb = colorbar('location', 'EastOutside'); set(get(cb,'ylabel'),'string','Mean BSR');
title('Theta 200 - 500 ms')

figureName = ['d2_erf_topo4'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');

%% 'Theta 1-2s'

h = figure('units','centimeters','position',[0 0 10 10]);
set(gcf,'renderer','Painters')

cfg = [];
cfg.layout = 'biosemi64.lay';
cfg.parameter = 'powspctrm';
cfg.comment = 'no';
cfg.colormap = cBrew;
cfg.colorbar = 'EastOutside';

plotData = [];
plotData.label = elec.label; % {1 x N}
plotData.dimord = 'chan';
cfg.zlim = [-4 4]; 
cfg.figure = h;
plotData.powspctrm = squeeze(nanmean(nanmean(stat.mask(:,stat.freq>2 &stat.freq<7,stat.time>1 & stat.time<1.5).*...
    stat.prob(:,stat.freq>2 & stat.freq<7,stat.time>1 & stat.time<1.5),3),2)); 
[~, sortind] = sort(plotData.powspctrm, 'ascend');
cfg.marker = 'off'; 
cfg.highlight = 'yes';
cfg.highlightchannel = plotData.label([38, 47]); % 46, 47, 38
cfg.highlightcolor = [1 1 1];
cfg.highlightsymbol = '.';
cfg.highlightsize = 30;
ft_topoplotER(cfg,plotData);
cb = colorbar('location', 'EastOutside'); set(get(cb,'ylabel'),'string','Mean BSR');
title('Theta 1-1.5s')

figureName = ['d2_erf_topo2'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');

%% title('AlphaBeta 0.5-1s')

h = figure('units','centimeters','position',[0 0 10 10]);
set(gcf,'renderer','Painters')

cfg = [];
cfg.layout = 'biosemi64.lay';
cfg.parameter = 'powspctrm';
cfg.comment = 'no';
cfg.colormap = cBrew;
cfg.colorbar = 'EastOutside';

plotData = [];
plotData.label = elec.label; % {1 x N}
plotData.dimord = 'chan';
cfg.zlim = [-1.5 1.5]; 
cfg.figure = h;
plotData.powspctrm = squeeze(nanmean(nanmean(stat.mask(:,stat.freq>8 & stat.freq<25,stat.time>0.6 & stat.time<1).*...
    stat.prob(:,stat.freq>8 & stat.freq<25,stat.time>0.6 & stat.time<1),3),2)); 
[~, sortind] = sort(plotData.powspctrm, 'ascend');
cfg.marker = 'off'; 
cfg.highlight = 'yes';
cfg.highlightchannel = plotData.label(sortind(1:3)); % 20, 21, 12
cfg.highlightcolor = [0 0 0];
cfg.highlightsymbol = '.';
cfg.highlightsize = 30;
ft_topoplotER(cfg,plotData);
cb = colorbar('location', 'EastOutside'); set(get(cb,'ylabel'),'string','Mean BSR');
title('AlphaBeta 0.6-1s')

figureName = ['d2_erf_topo3'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');

%% plot using raincloud plot

groupsizes=result.num_subj_lst;
conditions=lv_evt_list;
conds = {'hit'; 'miss'};
condData = []; uData = [];
for indGroup = 1
    if indGroup == 1
        relevantEntries = 1:groupsizes(1)*numel(conds);
    elseif indGroup == 2
        relevantEntries = groupsizes(1)*numel(conds)+1:...
             groupsizes(1)*numel(conds)+groupsizes(2)*numel(conds);
    end
    for indCond = 1:numel(conds)
        targetEntries = relevantEntries(conditions(relevantEntries)==indCond);        
        condData{indGroup}(indCond,:) = result.vsc(targetEntries,indLV);
        uData{indGroup}(indCond,:) = result.usc(targetEntries,indLV);
    end
end

cBrew(1,:) = 2.*[.3 .1 .1];
cBrew(2,:) = [.6 .6 .6];

idx_outlier = cell(1); idx_standard = cell(1);
for indGroup = 1
    dataToPlot = uData{indGroup}';
    % define outlier as lin. modulation that is more than three scaled median absolute deviations (MAD) away from the median
    X = [1 1; 1 2]; b=X\dataToPlot'; IndividualSlopes = b(2,:);
    outliers = isoutlier(IndividualSlopes, 'median');
    idx_outlier{indGroup} = find(outliers);
    idx_standard{indGroup} = find(outliers==0);
end

h = figure('units','centimeter','position',[0 0 8 10]);
ax = subplot(1,1,1);
for indGroup = 1
    dataToPlot = uData{indGroup}';
    % read into cell array of the appropriate dimensions
    data = []; data_ws = [];
    for i = 1:2
        for j = 1:1
            data{i, j} = dataToPlot(:,i);
            % individually demean for within-subject visualization
            data_ws{i, j} = dataToPlot(:,i)-...
                nanmean(dataToPlot(:,:),2)+...
                repmat(nanmean(nanmean(dataToPlot(:,:),2),1),size(dataToPlot(:,:),1),1);
            data_nooutlier{i, j} = data{i, j};
            data_nooutlier{i, j}(idx_outlier{indGroup}) = [];
            data_ws_nooutlier{i, j} = data_ws{i, j};
            data_ws_nooutlier{i, j}(idx_outlier{indGroup}) = [];
            % sort outliers to back in original data for improved plot overlap
            data_ws{i, j} = [data_ws{i, j}(idx_standard{indGroup}); data_ws{i, j}(idx_outlier{indGroup})];
        end
    end

    % IMPORTANT: plot individually centered estimates, stats on uncentered estimates!

    %subplot(1,2,indGroup);
    set(gcf,'renderer','Painters')
        cla;
        cl = cBrew(indGroup,:);
        [~, dist] = rm_raincloud_fixedSpacing(data_ws, [.8 .8 .8],1);
        h_rc = rm_raincloud_fixedSpacing(data_ws_nooutlier, cl,1,[],[],[],dist);
        view([90 -90]);
        axis ij
    box(gca,'off')
    set(gca, 'YTickLabels', {conds{2}; conds{1}});
    yticks = get(gca, 'ytick'); ylim([yticks(1)-(yticks(2)-yticks(1))./2, yticks(2)+(yticks(2)-yticks(1))./1.5]);

    minmax = [min(min(cat(2,data_ws{:}))), max(max(cat(2,data_ws{:})))];
    xlim(minmax+[-0.2*diff(minmax), 0.2*diff(minmax)])
    ylabel('recognition'); xlabel({'Brainscore'; '[Individually centered]'})

    % test linear effect
    curData = [data_nooutlier{1, 1}, data_nooutlier{2, 1}];
    X = [1 1; 1 2]; b=X\curData'; IndividualSlopes = b(2,:);
    [~, p, ci, stats] = ttest(IndividualSlopes);
    title(['M:', num2str(round(mean(IndividualSlopes),3)), '; p=', num2str(round(p,6))])
end
% make misses initial tick
ax.XDir = 'reverse';
set(ax, 'YTickLabels', {conds{1}; conds{2}});

figureName = ['d2_erf_rcp'];
saveas(h, fullfile(pn.figures, figureName), 'epsc');
saveas(h, fullfile(pn.figures, figureName), 'png');