function [] = figure7and8(correctionOff)

% function [] = figure7and8(correctionOff)
%
% without arguments, this plots figures 7 and 8 and part of figureS1 of
% Grenzebach, J., Wegner, T. G. G., Einhäuser, W., & Bendixen, A. (2021).
% Pupillometry in auditory multistability.
% PLoS ONE, 16(6), e0252370. https://doi.org/10.1371/journal.pone.0252370
%
% to obtain the uncorrected replay-active plot as shown in figure S3 of the
% supplement, call with argument correctionOff=1, i.e., figure7and8(1);

% by default, plot figure7 and 8
if nargin<1
    correctionOff=0;
end

maxData    = 2000; % data to be used to either side of 0, in ms
dataLength = 2*maxData+1;

conditionName = {'multistable','replay','random'};
plotColor     = {'b','r','g'};

% number of subjects
nSubs = 19;

% to collect data by Condition and subject
meanPupil   = nan(length(conditionName),nSubs,dataLength);
meanMissing = nan(length(conditionName),nSubs,dataLength);


for conditionCount = 1:length(conditionName)
    
    currentCond = conditionName{conditionCount};
    
    % load button press data
    load(sprintf('dataStudy2_buttons_%s',conditionName{conditionCount}),'buttonStates','buttonTimes');
    
    if strcmp(conditionName{conditionCount},'replay')
        
        % load stimulus change data (only needed for replay condition)
        load dataStudy2_stimPhases_replay stimPhaseStart
        
        % get the replay passive trace for correction
        [meanPupilRepPass,meanMissingRepPass]  = study2_replayPassive;
        
    end
    
    % load pupil data
    load(sprintf('dataStudy2_pupil_%s',currentCond),'pupilSize','xEye','yEye','tEye');
    
    
    % shorthands
    nSubjects = size(pupilSize,1);
    nBlocks   = size(pupilSize,2);
    
    % sanity check
    if nSubjects~=nSubs
        error('number of subjects in data file does not match expectation');
    end
    
    for subjectCount = 1 : nSubjects
        
        % to collect data per subject
        allPupil    = nan(0,dataLength);
        allMissing  = nan(0,dataLength);
        
        
        for blockCount = 1 : nBlocks
            
            % short hand notations for data of the current block
            currentPupilSize = pupilSize{subjectCount,blockCount};
            currentXEye      = xEye{subjectCount,blockCount};
            currentYEye      = yEye{subjectCount,blockCount};
            
            % remove data outside a central 10deg square
            currentPupilSize((currentXEye>1285)|(currentXEye<635))=nan;
            currentPupilSize((currentYEye>865) |(currentYEye<215))=nan;
            clear currentXEye currentYEye
            
            % remove the first 3s from analysis
            currentPupilSize(1:3000)=nan;
            
            
            % with this data, do the z-normalization
            pupilZ = (currentPupilSize-nanmean(currentPupilSize))/nanstd(currentPupilSize);
            clear currentPupilSize;
            
            % process the button presses
            currentButtonStates = buttonStates{subjectCount,blockCount};
            currentButtonTimes  = buttonTimes{subjectCount,blockCount};
            
            % find all state changes that correspond to a button press (not a
            % release), irrespective of the subsequent state
            idx = find(diff(currentButtonStates)>0)+1;
            
            % these are now defined as the transition times (i.e., switch times)
            transitionTimes = currentButtonTimes(idx);
            
            
            % for the replay condition, correction with the replay active
            % condition needs to be applied, and some button presses need to
            % be removed
            
            if strcmp(currentCond,'replay')
                
                % for the replay condition, some further conditions apply
                currentStimPhaseStart = stimPhaseStart{subjectCount,blockCount};
                removeTransitionIdx = [];
                
                % loop over all transition times
                for k = 1 : length(transitionTimes)
                    
                    % preceding stimulus change
                    ix = max(find(currentStimPhaseStart<transitionTimes(k)));
                    
                    if isempty(ix)
                        % there was no preceding stimulus change, remove transition
                        removeTransitionIdx(end+1) = k;
                    else
                        
                        alreadyRemoved = 0;
                        
                        if (k~=1)
                            
                            % this is not the first transition time
                            
                            if transitionTimes(k-1)>currentStimPhaseStart(ix)
                                
                                % the preceding transition followed the same stimulus
                                % change (only use the first button press after a
                                % stimulus change)
                                
                                removeTransitionIdx(end+1) = k;
                                alreadyRemoved=1;
                                
                            end
                            
                        end
                        
                        
                        if ((transitionTimes(k)-currentStimPhaseStart(ix))>2000)&(~alreadyRemoved)
                            
                            removeTransitionIdx(end+1) = k;
                            
                        end
                        
                    end
                    
                end
                
                % now remove the transitions that violate inclusion criteria
                transitionTimes(removeTransitionIdx)=[];
                
                % for the replay-passive subtraction, we in addition need the
                % index of the preceding stimulus change
                stimChangeTime = nan(1,length(transitionTimes));
                for k = 1:length(transitionTimes)
                    stimChangeTime(k) = currentStimPhaseStart(max(find(currentStimPhaseStart<transitionTimes(k))));
                end
                idxStartPassive = 1-(stimChangeTime-transitionTimes); % as the index corresponds to time in ms, this is fine
                idxEndPassive   = dataLength+1-(stimChangeTime-transitionTimes);
                
            end
            
            
            
            % code from here on again applies to all conditions
            
            
            % time of preceding and next transition
            prevTimes       = [0, transitionTimes(1:end-1)]; % at onset, this is trial start
            nextTimes       = [transitionTimes(2:end), currentButtonTimes(end)]; % at offset, this is trial end
            
            
            % define the usable time window for each transition
            startIdx = nan(1,length(transitionTimes));
            endIdx   = nan(1,length(transitionTimes));
            idxLeft = nan(1,length(transitionTimes));
            idxRight = nan(1,length(transitionTimes));
            
            for k = 1 : length(transitionTimes)
                
                % start and end of analysed epoch (-2s:2s)
                startIdx(k) = max(round((prevTimes(k)+transitionTimes(k))/2),transitionTimes(k)-maxData);
                endIdx(k)   = min(round((nextTimes(k)+transitionTimes(k))/2),transitionTimes(k)+maxData);
                
                % index in aggregated trace (1,4001 for full length, centred
                % around 2001 otherwise)
                idxLeft(k)  = 1+maxData+startIdx(k)-transitionTimes(k);
                idxRight(k) = endIdx(k)-transitionTimes(k)+maxData+1;
                
            end
            
            
            % now generate the traces
            localPupil   = nan(length(transitionTimes),dataLength);
            localMissing = nan(length(transitionTimes),dataLength);
            
            for k = 1 : length(transitionTimes)
                
                
                if strcmp(currentCond,'replay')&(~correctionOff)
                    
                    
                    tmp = meanPupilRepPass(subjectCount,idxStartPassive(k):idxEndPassive(k));
                    localPupil(k,idxLeft(k):idxRight(k))    = ...
                        pupilZ(startIdx(k):endIdx(k))-tmp(idxLeft(k):idxRight(k));
                else
                    localPupil(k,idxLeft(k):idxRight(k)) = ...
                        pupilZ(startIdx(k):endIdx(k));
                end
                
                
                localMissing(k,1:dataLength)            = isnan(localPupil(k,:));
            end
            
            % aggregate data across blocks
            allPupil    = [allPupil;localPupil];
            allMissing  = [allMissing;localMissing];
            
        end % blockCount
        
        % mean per subject and condition
        meanPupil(conditionCount,subjectCount,:)   = nanmean(allPupil);
        
        % only use traces for the computation of missing data points over
        % time that include at least one data point
        idxGood = find(sum(allMissing')~=dataLength);
        meanMissing(conditionCount,subjectCount,:) = nanmean(allMissing(idxGood,:),1);
        
        % baseline correction
        baseline = mean(meanPupil(conditionCount,subjectCount,1:200));
        meanPupil(conditionCount,subjectCount,:) = meanPupil(conditionCount,subjectCount,:) - baseline;
        
    end % subjectCount
    
end % conditionCount

% do the plotting
figure(7);
clf;
hold on;
legendHandle = nan(1,length(conditionName));

for conditionCount = 1 : length(conditionName)
    
    M   = squeeze(mean(meanPupil(conditionCount,:,:)));
    SEM = squeeze(std(meanPupil(conditionCount,:,:)))/sqrt(size(meanPupil,2));
    
    legendHandle(conditionCount) = ...
        plot([-maxData:maxData]/1000,M,plotColor{conditionCount});
    plot([-maxData:maxData]/1000,M-SEM,strcat(plotColor{conditionCount},':'));
    plot([-maxData:maxData]/1000,M+SEM,strcat(plotColor{conditionCount},':'));
    
end

ylabel('pupil size [z]');
xlabel('time to button press [s]');
axis([-2 2 -.1 1]);
if correctionOff
    ylim([-.1 1.2]);
    title('replay uncorrected (as in figureS3)');
end
set(gca,'XTick',[-2:1:2]);
set(gca,'YTick',[-.1:.1:1]);
line([-2 2],[0 0],'Color','k');
line([0 0],[-.1 1],'Color','k');
legend(legendHandle,conditionName);


%% jackknife [figure 8]
figure(8);
clf;
for conditionCount = 1 : numel(conditionName)
    
    currentPupil = squeeze(meanPupil(conditionCount,:,:));
    
    % allocate space for jackknifed traces
    jackknifedPupil = nan(size(currentPupil));
    
    
    for k = 1 : size(currentPupil,1)
        subs = setdiff([1:size(currentPupil,1)],k);
        jackknifedPupil(k,:) = mean(currentPupil(subs,:));
    end
    
    % plot all traces
    subplot(1,2,1);
    hold on;
    plot([-maxData:maxData]/1000,jackknifedPupil,plotColor{conditionCount});
    
    % plot maxima
    subplot(1,2,2);
    hold on;
    plot(conditionCount+[-nSubjects/2+.5:nSubjects/2-.5]*.02,max(jackknifedPupil'),strcat(plotColor{conditionCount},'o'));
    
end

subplot(1,2,1);
ylabel('pupil size [z]');
xlabel('time to button press [s]');
axis([-2 2 -.1 1]);
if correctionOff
    ylim([-.1 1.2]);
    title('replay uncorrected (as in figureS3)');
end
set(gca,'XTick',[-2:1:2]);
set(gca,'YTick',[-.1:.1:1]);
line([-2 2],[0 0],'Color','k');
line([0 0],[-.1 1],'Color','k');
legend(legendHandle,conditionName);


subplot(1,2,2);
ylabel('max pupil size [z]');
xlabel('condition');
axis([0.5 length(conditionName)+0.5 -.1 1]);
if correctionOff
    ylim([-.1 1.2]);
    title('replay uncorrected (as in figureS3)');
end
set(gca,'XTick',[1:1:length(conditionName)]);
set(gca,'XTickLabel',conditionName);
set(gca,'YTick',[-.1:.1:1]);
line([0.5 length(conditionName)+0.5],[0 0],'Color','k');


%% figureS1
figure(22);
clf;
for conditionCount = 1 : length(conditionName)
    subplot(4,1,conditionCount);
    M   = squeeze(mean(100*meanMissing(conditionCount,:,:)));
    SEM = squeeze(std(100*meanMissing(conditionCount,:,:)))/sqrt(size(meanPupil,2));
    
    plot([-maxData:maxData]/1000,M,plotColor{conditionCount});
    hold on;
    plot([-maxData:maxData]/1000,M-SEM,strcat(plotColor{conditionCount},':'));
    plot([-maxData:maxData]/1000,M+SEM,strcat(plotColor{conditionCount},':'));
    title(conditionName{conditionCount});
    ylabel('missng data [%]');
    xlabel('time to button press [s]');
    axis([-2 2 0 100]);
    set(gca,'XTick',[-2:1:2]);
    set(gca,'YTick',[0 50 100]);
    line([0 0],[0 100],'Color','k');
end

% same plot for replay passive
subplot(4,1,4);
M   = squeeze(mean(100*meanMissingRepPass(:,:)));
SEM = squeeze(std(100*meanMissingRepPass(:,:)))/sqrt(size(meanPupil,2));

plot([-2000:4000]/1000,M,'m');
hold on;
plot([-2000:4000]/1000,M-SEM,'m:');
plot([-2000:4000]/1000,M+SEM,'m:');
ylabel('missng data [%]');
xlabel('time to button press [s]');
axis([-2 4 0 100]);
set(gca,'XTick',[-2:1:5]);
set(gca,'YTick',[0 50 100]);
line([0 0],[0 100],'Color','k');
title('replay passive')
