% Two-armed bandit for the cognitive assessment task
% Version 2.0 - April 12, 2016
% C. Hassall
%    - 5 blocks of 20 trials, better instructions
% Version 3.0 - B. Norton and O. Krigolson, July 13, 2016
%    - Made stand-alone and for MUSE only
%    - Data recording is done using MATLAB; MUSELAB is not required for
%    this version.
% Version 3.1 - B. Norton, February 12, 2017
%   -Fixed baseline data
%   -Added debugging mode

% Block with summary - how many times they picked the correct one within that block (e.g. 14/20)
% Frame it as a game - your goal is to pick the best one on every trial
% (perfect score is 20 - hard to do)
% Block score
% Best score

clear all;
close all;
clc;
rng('shuffle');

debugging = 0; % Set to 1 to run experiments in small window for debugging purposes

data_file_name = input('Enter a subject number: ','s');

% If not in debugging mode, grab keyboard input and hide the mouse cursor
if ~debugging
    ListenChar(2);
    HideCursor();
end

%%% OPEN THE CONNECTION
osc_stream = osc_new_server(5555); % creates a new OSC server on Port 5555

%%% SET UP THE DRAWING WINDOW
background_colour = [0 0 0];
Screen('Preference','SkipSyncTests', 1) % ensures it will run, should be removed if possible
if ~debugging
    [win, rec] = Screen('OpenWindow', 0, 0); % Full screen
else
    [win, rec] = Screen('OpenWindow', 0, background_colour, [0 0 800 600],32,2); % Windowed, for debugging
end

%%% DEFINE VARIABLES
KbName('UnifyKeyNames');
p_data = [];
win_eeg = [];
lose_eeg = [];
eeg_fn = [data_file_name '_decision.mat'];
ExitKey = KbName('ESCAPE');

% Task parameters
num_blocks = 5; % Default: 5 block
num_trials = 20; % Default: 20 trials per block
p_win_1 = 0.6; % P(win) when square 1 is chosen
p_win_2 = 0.1; % P(win) when square 2 is chosen

% Task Variables
total_reward = 0;
best_score = 0;
num_wins = 0;
num_losses = 0;

% OSC Stuff
osc_wait_factor = 1.5; % This factor determines how long we'll wait for data, e.g. if we're trying to get 1000 ms of data then we are willing to wait up to 1500 ms
osc_timeout = 0.1; % Timeout for receiving OSC messages
length_of_baseline = 0.2;
length_of_erp = 1;

% * * * Graphical Properties * * *
normal_font_size = 20;
normal_font = 'Arial';

% Physical display properties
viewingDistance = 850; % mm, approximately
screenWidth = 598; % mm
screenHeight = 338; % mm
horizontalResolution = 1920; % Pixels
verticalResolution = 1080; % Pixels
horizontalPixelsPerMM = horizontalResolution/screenWidth;
verticalPixelsPerMM = verticalResolution/screenHeight;

% Screen/stim properties
squareDegrees = 2;
squareMMs = 2 * viewingDistance *tand(squareDegrees/2);
squareHorizontalPixels = horizontalPixelsPerMM * squareMMs;
squareVerticalPixels = verticalPixelsPerMM * squareMMs;
square_thickness = 30;
square_locations = [-squareHorizontalPixels*1.5 squareHorizontalPixels*1.5]; % Relative to center of the window

% Fixation
fixation_colour = [0 0 0];
go_colour = [180 180 180];
fixation_size = 60;

% Text
text_size = 40;
text_colour = [255 255 255];

xmid = rec(3)/2; ymid = rec(4)/2; % Based on the window that was opened
squares = [xmid-squareHorizontalPixels/2+square_locations(1) ymid-squareVerticalPixels/2 xmid+squareHorizontalPixels/2+square_locations(1) ymid+squareVerticalPixels/2;
    xmid-squareHorizontalPixels/2+square_locations(2) ymid-squareVerticalPixels/2 xmid+squareHorizontalPixels/2+square_locations(2) ymid+squareVerticalPixels/2;];

% Instructions (what the participant sees first)
instructions = 'In this game you see a series of coloured squares.\nOne of these coloured squares is the target - you are more likely to win and\nearn points if you select the target than if you select the other coloured square.\nYour goal is to figure out which coloured square is the target and select it as many times as possible.\n\nOn each trial of the game you will see a cross in the middle of the screen.\nTry to keep your eyes on the cross in the middle of the screen at all times.\nWhen the cross changes colour, pick the square that you think is the target by pressing either the f key or the j key.\nThroughout the game the squares will change colour. When this happens you will need to determine which one is the new target.\nRemember, your goal is to figure out which of the coloured squares is the target and to win as many times as possible.\n(Press space bar to continue.)'; 
questions = 'Any questions?\n(press space bar to proceed)';

Screen(win,'TextFont',normal_font);
Screen(win,'TextSize',normal_font_size);
DrawFormattedText(win, instructions,'center', 'center', text_colour,[],[],[],2);
Screen('Flip',win);

% Wait until keys are released
[keyIsDown, ~, ~, ~] = KbCheck;
while keyIsDown
    [keyIsDown, ~, ~, ~] = KbCheck;
end
WaitSecs(0.2);

% Wait for space bar press
done_looking = 0;
while ~done_looking
    [~, ~, keyCode, ~] = KbCheck;
    if strcmp(KbName(keyCode),'space') 
        done_looking = 1;
    end
end

DrawFormattedText(win, questions,'center', 'center', text_colour,[],[],[],2);
Screen('Flip',win);

% Wait until keys are released
[keyIsDown, ~, ~, ~] = KbCheck;
while keyIsDown
    [keyIsDown, ~, ~, ~] = KbCheck;
end
WaitSecs(0.2);

% Wait for space bar press
done_looking = 0;
while ~done_looking
    [~, ~, keyCode, ~] = KbCheck;
    if strcmp(KbName(keyCode),'space') 
        done_looking = 1;
    end
end

for b = 1:num_blocks
    
    % Pick two colours at random
    colours(1,:) = round(255 .* rand(1,3)); % Pick a random colour
    colours(2,:) = 255 - colours(1,:); % Use the complement
    
    % Remind the participant that this is a new block with new squares
    if b ~= 1
        Screen(win,'TextFont',normal_font);
        Screen(win,'TextSize',normal_font_size);
        DrawFormattedText(win, 'New Squares, New Target','center', 'center', text_colour,[],[],[],2);
        Screen('Flip',win);
        WaitSecs(1);
    end
    
    % How often participant chooses the correct square in this block
    performance = 0;
    
    % Trial loop
    for t = 1:num_trials
        
        % Draw crosshairs for 500 ms
        Screen('TextSize',win,fixation_size);
        DrawFormattedText(win, '+','center', 'center', fixation_colour,[],[],[],2);
        [~, ~, ~, ~, ~] = Screen('Flip',win);
        WaitSecs(0.5);
        
        % Randomize the order of square presentation
        if rand < 0.5
            order = [1 2];
        else
            order = [2 1];
        end
        
        % Draw the squares for 500 ms
        Screen('TextSize',win,fixation_size);
        DrawFormattedText(win, '+','center', 'center', fixation_colour,[],[],[],2);
        Screen('FillRect', win , colours(order(1),:), squares(1,:), square_thickness); % Left
        Screen('FillRect', win , colours(order(2),:), squares(2,:), square_thickness); % Right
        [~, ~, ~, ~, ~] = Screen('Flip',win);
                
        % Change the colour of the fixation cross ("go cue")
        Screen('TextSize',win,fixation_size);
        DrawFormattedText(win, '+','center', 'center', go_colour,[],[],[],2);
        Screen('FillRect', win , colours(order(1),:), squares(1,:), square_thickness); % Left
        Screen('FillRect', win , colours(order(2),:), squares(2,:), square_thickness); % Right
        [~, ~, ~, ~, ~] = Screen('Flip',win);
                
        % Wait for a response
        start_time = GetSecs();
        elapsed_time = 0;
        chosen_side = 0;
        while 1
            [keyIsDown, secs, keyCode, deltaSecs] = KbCheck;
            if strcmp(KbName(keyCode),'f')
                chosen_side = 1;
                elapsed_time = GetSecs() - start_time;
                break
            elseif strcmp(KbName(keyCode),'j')
                chosen_side = 2;
                elapsed_time = GetSecs() - start_time;
                break
            end
        end

        % Determine outcome
        this_trial_a_winner = -1;
        chosen_square = order(chosen_side); % [ 1 2 ]
        if chosen_square == 1
            this_trial_a_winner = rand < p_win_1;
            performance = performance + 1;
        elseif chosen_square == 2
            this_trial_a_winner = rand < p_win_2;
        end
              
        % Back to fixation
        Screen('TextSize',win,fixation_size);
        DrawFormattedText(win, '+','center', 'center', go_colour,[],[],[],2);
        [~, ~, ~, ~, ~] = Screen('Flip',win);
        
        % Clear the MUSE OSC buffer by calling osc_recv
        osc_recv(osc_stream,osc_timeout);
        
        jitter_amount = (rand() - 0.5) / 5;
        WaitSecs(0.3 + jitter_amount);
        
        % Record 200 ms of baseline
        num_samples = length_of_baseline * 500;
        temp_baseline = nan(4,num_samples);
        counter = 1;
        tic;
        while counter <= num_samples && toc < osc_wait_factor*length_of_baseline
            muse_data = osc_recv(osc_stream,osc_timeout);
            if ~isempty(muse_data)
                for p = 1:length(muse_data)
                    this_path = muse_data{p}.path;
                    if strcmp(this_path,'/muse/eeg')
                        temp_baseline(:,counter) = double(cell2mat(muse_data{p}.data))';
                        counter = counter + 1;
                        if counter > num_samples
                            break;
                        end
                    end
                end
            end
        end
        
        % Display the reward for 1 second
        Screen('TextSize',win,text_size);
        if this_trial_a_winner
            num_wins = num_wins + 1;
            total_reward = total_reward + 1;
            DrawFormattedText(win, 'WIN','center', 'center', text_colour,[],[],[],2);
            [~, ~, ~, ~, ~] = Screen('Flip',win);
        else
            num_losses = num_losses + 1;
            DrawFormattedText(win, 'LOSE','center', 'center', text_colour,[],[],[],2);
            [~, ~, ~, ~, ~] = Screen('Flip',win);
        end
        
        % Record 800 ms of EEG
        num_samples = length_of_erp * 500;
        temp_eeg = nan(4,num_samples);
        counter = 1;
        tic;
        while counter <= num_samples && toc <= osc_wait_factor*length_of_erp
            muse_data = osc_recv(osc_stream,osc_timeout);
            if ~isempty(muse_data)
                for p = 1:length(muse_data)
                    this_path = muse_data{p}.path;
                    if strcmp(this_path,'/muse/eeg')
                        temp_eeg(:,counter) = double(cell2mat(muse_data{p}.data))';
                        counter = counter + 1;
                        if counter > num_samples
                            break;
                        end
                    end
                end
            end
        end
        
        WaitSecs(0.2); % Feedback is up for 1000 ms total
        
        if this_trial_a_winner
            this_eeg = [temp_baseline temp_eeg];
            win_eeg = [win_eeg this_eeg];
        else
            this_eeg = [temp_baseline temp_eeg];
            lose_eeg = [lose_eeg this_eeg];
        end
        
        % Save this trial
        EEG{1} = win_eeg;
        EEG{2} = lose_eeg;
        save(eeg_fn, 'EEG');
                
        this_data_line = [t colours(1,:) colours(2,:) order chosen_side chosen_square elapsed_time*1000 this_trial_a_winner performance];
        dlmwrite([data_file_name '_decision_behavioural'],this_data_line,'delimiter', '\t', '-append');
        
        [~, ~, keyCode] = KbCheck();
        if keyCode(ExitKey)
            break;
        end
        
    end
    
    % Keep track of best block
    if performance > best_score
        best_score = performance;
    end
    
    % Display block summary
    block_message = ['You picked the correct square ' num2str(performance) ' out of ' num2str(num_trials) ' times!\nBlock Score: ' num2str(performance) '\nBest Score: ' num2str(best_score)];
    Screen(win,'TextFont',normal_font);
    Screen(win,'TextSize',normal_font_size);
    DrawFormattedText(win, block_message, 'center', 'center', text_colour,[],[],[],2);
    Screen('Flip',win);
    [~, ~, keyCode] = KbCheck();
    if keyCode(ExitKey)
        break;
    end
    WaitSecs(2);
end

%%% CLOSE THE CONNECTION
osc_free_server(osc_stream); % releases the connection with MUSE

% Close the Psychtoolbox window and bring back the cursor and keyboard
Screen('CloseAll');
ListenChar(0);
ShowCursor();