clear all; close all;

Screen('Preference', 'SkipSyncTests', 1);

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

%% Set up general parameters

% ask for block number of staircase block
prompt = {'Enter block number:'};
dlg_title = 'Block Information';
num_lines = 1;
answer = inputdlg(prompt,dlg_title,num_lines);

Gparams                 = [];
Gparams.block_number    = str2num(answer{1});

Gparams = setupGParams(Gparams);

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

%% Initialize Quest or load staircase if block > 1
if Gparams.block_number == 1
    
    tGuess      = log10(2);
    tGuessSd    = 3;
    pThreshold  = 0.70;
    beta        = 3.5;
    delta       = 0.01;
    gamma       = 0.5;
    
    q = QuestCreate(tGuess,tGuessSd,pThreshold,beta,delta,gamma);
    q.normalizePdf=1; % This adds a few ms per call to QuestUpdate, but otherwise the pdf will underflow after about 1000 trials.
    
else
    
    filename = [Gparams.DataPath 'staircase_data_subj-' num2str(Gparams.subjectNumber) '_block-' num2str(Gparams.block_number-1) '.mat'];
    tmp = load(filename);
    q = tmp.q;
    
    clear tmp;
    
    % potentially better to retrain a new staircase with the last X trials
    % of the old staircase. When running more blocks, forgetting about
    % older trials will help against non-stationarity
    
end

nTrials = 44; % How many trials to run in one block of the staircasing
nEasyTrials = 4; % How many of nTrials are easy "starter" trials
iBlock = 1; % For the sake of compatibility with other scripts

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

% Directions of the standard and test stimuli
trialMatrix.direction(1,:) = zeros(1,nTrials);
trialMatrix.direction(1,1) = rand(1)*360;

for iTrial = 2:nTrials
    trialMatrix.direction(1,iTrial) = mod(trialMatrix.direction(1,iTrial-1) + (rand(1)*300+30),360); % Each trial is at least 30 degrees offset from the previous trial direction
end

trialMatrix.responseMapping(1,:) = ones(1,nTrials) * (mod(Gparams.subjectNumber,2) + 1);   % Response mapping depends on subject number

% relative coherence will change from trial to trial depending on the
% subject's performance, so it cannot be prespecified here

% vector to specify increment or decrement in coherence
increment = [-1 * ones(1,nTrials/2), ones(1,nTrials/2)];
perm_vec  = randperm(nTrials);
increment = increment(perm_vec);


%% Allocate space for flip and response data

results                 = [];

results.responseTime    = NaN(1,nTrials);
results.responseKey     = NaN(1,nTrials);
results.responseDir     = NaN(1,nTrials);
results.responseCorrect = NaN(1,nTrials);

flip                    = [];

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

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

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

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

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

%% Instructions

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

% Instruction screen
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' ...
    'You have only ' num2str(Gparams.timing.responseTimeDirectionOnly) ' seconds to respond.\n\n'...
    'Try to answer as quickly and accurately as possible.\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 response.\n\n\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.'];
showtext(text,Gparams);
    
text = ['START PRACTICE\n\n\n\n' ...
    'Now you will practice the task again.\n\n\n\n'...
    'The practice will start with some rather easy rounds\n\n' ...
    'and will become more difficult over time.\n\n\n\n'...
    'At the start of the round you will see a reminder\n\n' ...
    'of which button ("j" and "k") to use for which response.\n\n' ...
    'After that, the task will follow automatically.\n\n\n\n' ...
    'Press "space" to start 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:nTrials
    
    %% 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);
    
    
    %% Figure out what stimulus coherence to show and create the stimuli
    if iTrial > nEasyTrials
        stim_level = 10^(QuestQuantile(q));
        
        
        if stim_level > .20
            stim_level = .20;
        end
    else
        stim_level = .20;
    end
    
    trialMatrix.relative_coherence(1,iTrial) = stim_level * increment(iTrial);
    
    testStimCoord       = dots(iTrial,1,trialMatrix,0,Gparams);
    standardStimCoord   = dots(iTrial,1,trialMatrix,1,Gparams);
    
    
    %% 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(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;
    
    % 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(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();
        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.responseTimeDirectionOnly
        
        if ~keyIsDown1
            [keyIsDown1, secs1, keyCode1]  = 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
    
    % Compute correctness of response
    if isnan(results.responseDir(iBlock,iTrial))
        results.responseCorrect(iBlock,iTrial) = 0;
    elseif results.responseDir(iBlock,iTrial) == sign(trialMatrix.relative_coherence(iBlock,iTrial))
        results.responseCorrect(iBlock,iTrial) = 1;
    else
        results.responseCorrect(iBlock,iTrial) = 0;
    end
    
    %% Update Quest with the current response
    if iTrial > nEasyTrials
        q = QuestUpdate(q,log10(stim_level),results.responseCorrect(iBlock,iTrial));
    end
    
    %% Escape trial loop
    if Gparams.cancel == true
        break
        warning('experiment was manually terminated');
    end
    
    %% Feedback
    
    if ~keyIsDown1
        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

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

%% End of block message
RestrictKeysForKbCheck([KbName('Space')]);
    
text = ['This is the end of the round.\n\n\n' ...
    'Press "space" to close the screen\n\n\n' ...
    'Please report to the experimenter.'];
showtext(text,Gparams);

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

%% Diagnostics
nTrials2 = 38;

% Threshold, performance
figure; subplot(1,2,1);
plot(q.trialCount-nTrials2:q.trialCount,10.^(q.intensity(q.trialCount-nTrials2:q.trialCount)));
ylim([min(10.^(q.intensity(q.trialCount-nTrials2:q.trialCount))) - 1,  max(10.^(q.intensity(q.trialCount-nTrials2:q.trialCount))) + 1])

xlabel('Trial Number')
ylabel('Relative Coherence')

ylim([.0 .30])

threshold = 10^(QuestMean(q));
disp(['Threshold: ' num2str(threshold) ' +- ' num2str(10^(QuestSd(q)))]);

% Response time
subplot(1,2,2);
plot(results.responseTime);
xlim([0, nTrials]);
ylim([0, 4])

xlabel('Trial Number')
ylabel('Response time')

avgResponseTime = round(nanmean(results.responseTime),3);
disp(['Mean response time: ' num2str(avgResponseTime) ' seconds']);

% Missed trials
missed = sum(isnan(results.responseDir));
disp(['Missed trials: ' num2str(missed) ' out of ' num2str(length(results.responseDir))]);


%% Done staircasing? Set the difficulty level
prompt = 'Set difficulty? y/n [n]: ';
str = input(prompt,'s');
if isempty(str)
    str = 'n';
end

if strcmp(str,'y')
    
    if threshold <= 0.05
        difficulty = 3;
    elseif threshold <= 0.10
        difficulty = 2;
    else
        difficulty = 1;
    end
        
    save([Gparams.DataPath 'difficulty.mat'], 'difficulty');
end