clear all; close all;

rng('shuffle'); % Randomize seed of random number generator

%% Set up general parameters
Gparams = [];
Gparams = setupGParams(Gparams);

%% Hide Cursor and set priority
HideCursor;
priorityLevel = MaxPriority(Gparams.pWindow);
Priority(priorityLevel);

%% Generate block and trial structure
nBlocks         = 1;
nRepetitions    = 1;
cbOrder         = 1;

feedback        = true; % indicate whether participants will receive feedback about their decision

filename = [Gparams.DataPath 'difficulty.mat'];
tmp = load(filename);
difficulty = tmp.difficulty;

clear tmp

% Use the stimulus set with the difficulty indicated by the staircase
if difficulty == 1
    stimulus_set = Gparams.dots.relative_coherence_easy;
elseif difficulty == 2
    stimulus_set = Gparams.dots.relative_coherence_medium;
elseif difficulty == 3
    stimulus_set = Gparams.dots.relative_coherence_hard;
else 
    stimulus_set = Gparams.dots.relative_coherence;
end

for iBlock = 1:nBlocks
    
    % Relative coherences of the standard and test stimuli
    coh = repmat(stimulus_set,1,4)';
    trialMatrix.relative_coherence(iBlock,:) = coh(randperm(length(coh)));
    
    % Directions of the standard and test stimuli
    trialMatrix.direction(iBlock,:) = zeros(1,length(coh));
    trialMatrix.direction(iBlock,1) = rand(1)*360;
    
    for iTrial = 2:length(coh)
        trialMatrix.direction(iBlock,iTrial) = mod(trialMatrix.direction(iBlock,iTrial-1) + (rand(1)*300+30),360); % Each trial is at least 30 degrees offset from the previous trial direction
    end
    
    trialMatrix.responseMapping(iBlock,:) = ones(size(trialMatrix.relative_coherence(iBlock,:))) * (mod(Gparams.subjectNumber,2) + 1);   % Response mapping depends on subject number
      
end

if size(trialMatrix.relative_coherence) ~= size(trialMatrix.responseMapping)
   error('Trialsequence dimensions are not equal')   
end

% Save the trialmatrix
save([Gparams.DataPath 'refresher_trialmatrix_subj-' num2str(Gparams.subjectNumber) '_ses-' num2str(Gparams.sessionNumber) '_' datestr(now, 'yyyy-mm-dd_HH-MM-SS') '.mat'],'trialMatrix','-v7.3');

nTrials = length(trialMatrix.relative_coherence);

%% Allocate space for flip and response data

results                 = [];

results.responseTime    = NaN(nBlocks,nTrials);
results.responseKey     = NaN(nBlocks,nTrials);
results.responseDir     = NaN(nBlocks,nTrials);
results.responseCorrect = NaN(nBlocks,nTrials);
results.confidenceTime  = NaN(nBlocks,nTrials);
results.confidence      = NaN(nBlocks,nTrials);

flip                    = [];

flip.itiOnset.VBL       = NaN(nBlocks,nTrials);
flip.itiOnset.StimOnset = NaN(nBlocks,nTrials);
flip.itiOnset.FlipTS    = NaN(nBlocks,nTrials);
flip.itiOnset.Missed    = NaN(nBlocks,nTrials);
flip.itiOnset.BeamPos   = NaN(nBlocks,nTrials);

flip.standardStim.VBL       = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.standardStim.StimOnset = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.standardStim.FlipTS    = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.standardStim.Missed    = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.standardStim.BeamPos   = NaN(nBlocks,nTrials,Gparams.timing.nframes);

flip.isiOnset.VBL         = NaN(nBlocks,nTrials);
flip.isiOnset.StimOnset   = NaN(nBlocks,nTrials);
flip.isiOnset.FlipTS      = NaN(nBlocks,nTrials);
flip.isiOnset.Missed      = NaN(nBlocks,nTrials);
flip.isiOnset.BeamPos     = NaN(nBlocks,nTrials);

flip.testStim.VBL       = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.testStim.StimOnset = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.testStim.FlipTS    = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.testStim.Missed    = NaN(nBlocks,nTrials,Gparams.timing.nframes);
flip.testStim.BeamPos   = NaN(nBlocks,nTrials,Gparams.timing.nframes);

flip.testStimOffset.VBL       = NaN(nBlocks,nTrials);
flip.testStimOffset.StimOnset = NaN(nBlocks,nTrials);
flip.testStimOffset.FlipTS    = NaN(nBlocks,nTrials);
flip.testStimOffset.Missed    = NaN(nBlocks,nTrials);
flip.testStimOffset.BeamPos   = NaN(nBlocks,nTrials);

%% Instructions

RestrictKeysForKbCheck([KbName('Space')]);

text = ['TASK SUMMARY\n\n\n\n' ...
    'In the task you will always see two stimuli:\n\n' ...
    'a standard stimulus, followed by a second stimulus.\n\n' ...
    'You are asked to respond to the coherence difference,\n\n' ...
    'using the "j" and "k" keys.\n\n' ...
    'Immediately after, you will indicate your confidence\n\n' ...
    'about your decision using the "1", "2", "3" and "4" keys.\n\n' ...
    'You have ' num2str(Gparams.timing.responseTime) ' seconds total to do this.\n\n'...
    'Tip: you can respond as soon as the second stimulus appears.\n\n\n\n' ...
    'You will get feedback about the correctness of your decision.\n\n'...
    'Please remember to always maintain fixation on the cross\n\n' ...
    'in the center of the screen.\n\n\n\n' ...
    'Press "space" to continue.'];


%% run blocks
for iBlock = 1:nBlocks
    
    DrawFormattedText(Gparams.pWindow, 'Loading stimuli...', 'center', 'center', 255);
    Screen('Flip', Gparams.pWindow);
    
    % preload all the dot coordinates - overwrite this block!
    testStimCoord = nan(nTrials, Gparams.timing.nframes, 2, Gparams.dots.nDots);
    standardStimCoord = nan(nTrials, Gparams.timing.nframes, 2, Gparams.dots.nDots);
    
    for iTrial = 1:nTrials
        testStimCoord(iTrial, :, :, :) = dots(iTrial,iBlock,trialMatrix,0,Gparams);
        standardStimCoord(iTrial, :, :, :) = dots(iTrial,iBlock,trialMatrix,1,Gparams);
    end
    
    % Save the coordinates
    save([Gparams.DataPath 'dots_subj-' num2str(Gparams.subjectNumber) '_ses-' num2str(Gparams.sessionNumber) '_block-' num2str(iBlock) '_' datestr(now, 'yyyy-mm-dd_HH-MM-SS') '.mat'],'testStimCoord','-v7.3');
    save([Gparams.DataPath 'dots_subj-' num2str(Gparams.subjectNumber) '_ses-' num2str(Gparams.sessionNumber) '_block-' num2str(iBlock) '_' datestr(now, 'yyyy-mm-dd_HH-MM-SS') '.mat'],'standardStimCoord','-v7.3');
    
    % Instruction screen
    RestrictKeysForKbCheck([KbName('Space')]);
    
    text = ['START REFRESHER\n\n\n\n' ...
        'You will now do a brief round to refresh the task.\n\n\n\n' ...
        'At the start of the round you will see a reminder \n\n' ...
        'of which button ("j" or "k") you will use for which directional answer. \n\n' ...
        'After that, the task will follow automatically.\n\n\n\n' ...
        'Press "space" to start with the round.\n\n' ...
        'The round takes approximately 5 minutes.'];
    showtext(text,Gparams);

    % Show response mapping for this block
    if mean(trialMatrix.responseMapping(iBlock,:)) == 1
        DrawFormattedText(Gparams.pWindow, 'LOWER (j)                     HIGHER (k)', 'center', 'center', 255);
    elseif mean(trialMatrix.responseMapping(iBlock,:)) == 2
        DrawFormattedText(Gparams.pWindow, 'HIGHER (j)                     LOWER (k)', 'center', 'center', 255);       
    end
        drawFixation(Gparams);
        Screen('Flip', Gparams.pWindow);
        WaitSecs(3);

    
    for iTrial = 1:length(trialMatrix.relative_coherence)
        
        %% Fixation
        drawFixation(Gparams);
        [flip.itiOnset.VBL(iBlock,iTrial),...
            flip.itiOnset.StimOnset(iBlock,iTrial),...
            flip.itiOnset.FlipTS(iBlock,iTrial),...
            flip.itiOnset.Missed(iBlock,iTrial),...
            flip.itiOnset.BeamPos(iBlock,iTrial)] = Screen('Flip', Gparams.pWindow);
        
        
        %% Standard stimulus
        presentationTime = flip.itiOnset.VBL(iBlock,iTrial) + Gparams.timing.itiDirectionOnly - Gparams.frameDuration - Gparams.flipLag;
        % subtract frameDuration to compensate for frame spent between end
        % and beginning of for-loop (over trials)
        
        for iFrame = 1:Gparams.timing.nframes
            
            drawDots(squeeze(standardStimCoord(iTrial,iFrame, :, :)),Gparams);
            drawFixation(Gparams);
            [flip.standardStim.VBL(iBlock,iTrial,iFrame),...
                flip.standardStim.StimOnset(iBlock,iTrial,iFrame),...
                flip.standardStim.FlipTS(iBlock,iTrial,iFrame),...
                flip.standardStim.Missed(iBlock,iTrial,iFrame),...
                flip.standardStim.BeamPos(iBlock,iTrial,iFrame)] = ...
                Screen('Flip', Gparams.pWindow,presentationTime);
            
            presentationTime = flip.standardStim.VBL(iBlock,iTrial,iFrame) + Gparams.frameDuration-Gparams.flipLag;
            
        end
        
        %% ISI
        drawFixation(Gparams);
        [flip.isiOnset.VBL(iBlock,iTrial),...
            flip.isiOnset.StimOnset(iBlock,iTrial),...
            flip.isiOnset.FlipTS(iBlock,iTrial),...
            flip.isiOnset.Missed(iBlock,iTrial),...
            flip.isiOnset.BeamPos(iBlock,iTrial)] = Screen('Flip', Gparams.pWindow,presentationTime);
        
        
        %% Test Stimulus
        keyIsDown1 = 0;
        keyIsDown2 = 0;
        
        % Only allow direction buttons and escape
        RestrictKeysForKbCheck([Gparams.responseKeys.left ...
            Gparams.responseKeys.right ...
            Gparams.responseKeys.escape]);
        
        presentationTime = flip.isiOnset.VBL(iBlock,iTrial) + Gparams.timing.isi - Gparams.flipLag;
        
        for iFrame = 1:Gparams.timing.nframes
            
            drawDots(squeeze(testStimCoord(iTrial,iFrame, :, :)),Gparams);
            drawFixation(Gparams);
            [flip.testStim.VBL(iBlock,iTrial,iFrame),...
                flip.testStim.StimOnset(iBlock,iTrial,iFrame),...
                flip.testStim.FlipTS(iBlock,iTrial,iFrame),...
                flip.testStim.Missed(iBlock,iTrial,iFrame),...
                flip.testStim.BeamPos(iBlock,iTrial,iFrame)] = ...
                Screen('Flip', Gparams.pWindow,presentationTime);
            
            presentationTime = flip.testStim.VBL(iBlock,iTrial,iFrame) + Gparams.frameDuration-Gparams.flipLag;
            
            if ~keyIsDown1
                [keyIsDown1, secs1, keyCode1]  = KbCheck();
                if keyIsDown1
                    % Only allow confidence buttons and escape
                    RestrictKeysForKbCheck([Gparams.responseKeys.conf1...
                        Gparams.responseKeys.conf2...
                        Gparams.responseKeys.conf3...
                        Gparams.responseKeys.conf4...
                        Gparams.responseKeys.escape]);
                end
            elseif ~keyIsDown2 
                [keyIsDown2, secs2, keyCode2]  = KbCheck();
            end
             
        end
        
        %% Response (continued)
        
        % Draw fixation again
        drawFixation(Gparams);
        [flip.testStimOffset.VBL(iBlock,iTrial),...
            flip.testStimOffset.StimOnset(iBlock,iTrial),...
            flip.testStimOffset.FlipTS(iBlock,iTrial),...
            flip.testStimOffset.Missed(iBlock,iTrial),...
            flip.testStimOffset.BeamPos(iBlock,iTrial)] = ...
            Screen('Flip', Gparams.pWindow, presentationTime);
        
        % Allow people to still respond
        
        iTime = GetSecs;
        
        while iTime < presentationTime + Gparams.timing.responseTime
            
             if ~keyIsDown1
                [keyIsDown1, secs1, keyCode1]  = KbCheck();
                if keyIsDown1
                    % Only allow confidence buttons and escape
                    RestrictKeysForKbCheck([Gparams.responseKeys.conf1...
                        Gparams.responseKeys.conf2...
                        Gparams.responseKeys.conf3...
                        Gparams.responseKeys.conf4...
                        Gparams.responseKeys.escape]);
                end
            elseif ~keyIsDown2 
                [keyIsDown2, secs2, keyCode2]  = KbCheck();
            end
            
            iTime = GetSecs;
            
        end
                
        % Write the response into results
        
        if keyIsDown1
            
            % Log RT
            results.responseTime(iBlock, iTrial) = secs1 - flip.testStim.VBL(iBlock,iTrial,1);  % time from stimulus onset
            
            % Log button press
            % If both buttons are pressed simultaneously, record as invalid
            % response
            if length(find(keyCode1)) > 1
                results.responseKey(iBlock, iTrial) = NaN;
            else  
                switch KbName(find(keyCode1))
                    case 'j'
                        results.responseKey(iBlock, iTrial) = -1;
                    case 'k'
                        results.responseKey(iBlock, iTrial) = 1;
                    case 'esc'
                        Gparams.cancel = true;
                        results.responseKey(iBlock, iTrial) = NaN;
                    case 'ESCAPE' % if escape is pressed, exit the experiment
                        Gparams.cancel = true;
                        results.responseKey(iBlock, iTrial) = NaN;
                end
                
            end
            
            % Log response direction
            if ~isnan(results.responseKey(iBlock, iTrial))
               
                switch trialMatrix.responseMapping(iBlock,iTrial)
                    case 1
                        results.responseDir(iBlock,iTrial) = ...
                            results.responseKey(iBlock, iTrial);
                    case 2
                        results.responseDir(iBlock,iTrial) = ...
                            results.responseKey(iBlock, iTrial) * -1;
                end
                
            end
              
        end
        
        if keyIsDown2

            % Log RT
            results.confidenceTime(iBlock, iTrial) = secs2 - secs1;

            % Log button press
            % If both buttons are pressed simultaneously, record as invalid
            % response
            if length(find(keyCode2)) > 1
                results.confidence(iBlock, iTrial) = NaN;
            else  
                switch KbName(find(keyCode2))
                    case '1!'
                        results.confidence(iBlock, iTrial) = 1;
                    case '2@'
                        results.confidence(iBlock, iTrial) = 2;
                    case '3#'
                        results.confidence(iBlock, iTrial) = 3;
                    case '4$'
                        results.confidence(iBlock, iTrial) = 4;
                    case 'esc'
                        Gparams.cancel = true;
                        results.confidence(iBlock, iTrial) = NaN;
                    case 'ESCAPE' % if escape is pressed, exit the experiment
                        Gparams.cancel = true;
                        results.confidence(iBlock, iTrial) = NaN;
                end

            end
            
        end
        
        % Compute correctness of response
        if isnan(results.responseDir(iBlock,iTrial))
            results.responseCorrect(iBlock,iTrial) = NaN; 
        elseif results.responseDir(iBlock,iTrial) == sign(trialMatrix.relative_coherence(iBlock,iTrial))
            results.responseCorrect(iBlock,iTrial) = true;
        else
            results.responseCorrect(iBlock,iTrial) = false;    
        end    
            
        %% Escape trial loop
        if Gparams.cancel == true
            break
            warning('experiment was manually terminated');      
        end
        
        %% Feedback
        
        if ~keyIsDown1 || ~keyIsDown2
            PsychPortAudio('SetLoop',Gparams.audio.h, Gparams.sound.tonepos(3,1), Gparams.sound.tonepos(3,2));
        elseif results.responseCorrect(iBlock,iTrial) == true
            PsychPortAudio('SetLoop',Gparams.audio.h, Gparams.sound.tonepos(1,1), Gparams.sound.tonepos(1,2));
        else 
            PsychPortAudio('SetLoop',Gparams.audio.h, Gparams.sound.tonepos(2,1), Gparams.sound.tonepos(2,2));
        end    
        PsychPortAudio('Start', Gparams.audio.h);
        
    end
          
    %% Escape block loop
    if Gparams.cancel == true
            break
            warning('experiment was manually terminated');      
    end
                    
    %% End of block message    
    RestrictKeysForKbCheck([KbName('Space')]);
    
    if (iBlock < nBlocks)
        showtext(sprintf('You just finished round %g of %g.\n\n\n You had %.2f percent correct\n\nand you responded on average after %.2f seconds.\n\n You missed %g responses.\n\n\n\nFeel free to take a short break.\n\n\nWhen you are ready press "space" to continue...', ...
            iBlock, nBlocks, nanmean(results.responseCorrect(iBlock,:))*100, nanmean(results.responseTime(iBlock,:)), sum(isnan(results.responseDir(iBlock,:)))), Gparams);
    elseif iBlock == nBlocks
        showtext(sprintf('You have finished the experiment.\n\n\n In this last round you had %.2f percent correct\n\nand you responded on average after %.2f seconds.\n\n You missed %g responses.\n\n\n\n Please report to the experimenter.', ...
            nanmean(results.responseCorrect(iBlock,:))*100, nanmean(results.responseTime(iBlock,:)), sum(isnan(results.responseDir(iBlock,:)))), Gparams);
    end
    
    
    %% Print stats
    fprintf('\nBLOCK %g\n\nPERFORMANCE\n\n%.2f percent correct\nRT = %.2f s\nMissed responses: %d\n', ...
    iBlock,...
    nanmean(results.responseCorrect(iBlock,:))*100, nanmean(results.responseTime(iBlock,:)),sum(isnan(results.responseTime(iBlock,:))));
 
    fprintf('\nCorrect per difficulty level: %.2f, %.2f, %.2f, %.2f, %.2f, %.2f\n', ...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(6))), ...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(5))), ...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(4))),...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(3))), ...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(2))), ...
        100* nanmean(results.responseCorrect(iBlock, abs(trialMatrix.relative_coherence(iBlock,:))==stimulus_set(1))));
    
    fprintf('\nCONFIDENCE\n1 = %gx, 2 = %gx, 3 = %gx, 4 = %gx\n', ...
    size(results.confidence(results.confidence(iBlock,:)==1),2), size(results.confidence(results.confidence(iBlock,:)==2),2), size(results.confidence(results.confidence(iBlock,:)==3),2), size(results.confidence(results.confidence(iBlock,:)==4),2));

end

%% Save data of the experiment
datetime = datestr(now, 'yyyy-mm-dd_HH-MM-SS');
filename = [Gparams.DataPath 'refresher_data_subj-' num2str(Gparams.subjectNumber) '_ses-' num2str(Gparams.sessionNumber) '_' datetime '.mat'];
save(filename,'Gparams','trialMatrix','results','flip');

%% Cleaning up
Screen('CloseAll');
ShowCursor;
ListenChar(0);
Priority(0);