% P300 Task
% Version 1.0
% The Cog Assess oddball task, credit to Cam Hassall for the original, Bryn
% and Chad for work on this, some help from Olav Krigolson
% version

% Muse preset settings to use
% muse-io --no-dsp --preset AD --osc osc.udp://localhost:5555 --device Muse-XXXX

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% SETUP COMMANDS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clc;                                % clear the matlab console
close all;                          % close all open elements
clear all;                          % clear the memory
rng('shuffle');                     % seed the random number generator

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% VARIABLES
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

sample_rate = 500;                  % Muse sample rate, in Hz
length_of_baseline = 0.2;           % length of the ERP baseline to be recorded, in seconds
length_of_erp = 0.8;                % length of the ERP window to be recorded, in seconds
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

number_of_blocks = 2;               % number of blocks of trials
trials_per_block = 10;              % number of trials in a block
oddball_probability = 0.25;         % probability that any given trial is an oddball

background_colour = [0 0 0];        % baclground colour in RGB coordinates

normal_font_size = 20;              % text font size
normal_font = 'Arial';              % text font
text_colour = [255 255 255];        % text colour

fixation_size = 20;                 % size of the fixation cross
fixation_colour = [255 255 255];    % fixation colour

stim_size = 20;                     % size of fixation stimulus
oddball_colour = [0 0 255];         % oddball colour
control_colour = [0 255 0];         % control colour

ExitKey = KbName('ESCAPE');         % key to break out of experimental loop

osc_port = 5555;                    % port that incoming OSC stream is on

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% CHANGE BELOW HERE AT YOUR OWN RISK
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Get participant info
while 1
    clc;
    p_number = input('Enter the participant number:\n','s');  % get the subject name/number
    datafile_name = strcat(['Oddball_' p_number '.mat']);
    checker = exist(datafile_name,'file');
    if checker == 0
        break;
    else
        disp('Filename Already Exists!');
        WaitSecs(1);
    end
end
behavioural_file_name = strcat(['Oddball_Behaviour_' p_number]);

% Hide the cursor and "disable" keyboard
% Press ctl-c to bring it back
ListenChar(2);
HideCursor();

% add the osc stream if it does not already exist
if ~exist('osc_stream','var')
    osc_stream = osc_new_server(osc_port);
end

% set up some variables
odd_eeg = [];   % where the odd trial eeg is stored
dist_eeg = [];   % where the control trial eeg is stored
data = [];    % where the behavioural data is stored
ALLEEG = []; % concatenated odd and control eeg
markers = []; % List of markers
num_odd = 0; % oddball counter
num_dist = 0; % control counter
actual_trial_counter = 1; % trial counter (all trials)

try
    
    % Set up a PTB window
    Screen('Preference','SkipSyncTests', 1); % This may need to be uncommented on some machines
    [win, rec] = Screen('OpenWindow', 0, background_colour, [],32,2); % use [0 0 800 600] for testing, [] for running
    % Calculate the center of the screen given the screen size
    xmid = rec(3)/2;
    ymid = rec(4)/2;
    
    % Display welcome screen
    summary = 'Welcome. Press any key when you are ready to begin.';
    Screen(win,'TextColor',[255 255 255]);
    Screen(win,'TextFont',normal_font);
    Screen(win,'TextSize',normal_font_size);
    DrawFormattedText(win, summary,'center', 'center', text_colour,[],[],[],2);
    Screen('Flip',win);
    KbStrokeWait();

    % Display oddball instructions
    % Instructions (what the participant sees first)
    instructions = 'You are going to see a series of blue or green circles in the middle of the screen\nThe blue circle is the target. Count the number of blue circles you see\nTry to keep your eyes on the middle of the display at all times\n(press any key to proceed)';
    Screen(win,'TextFont',normal_font);
    Screen(win,'TextSize',normal_font_size);
    DrawFormattedText(win, instructions,'center', 'center', text_colour,[],[],[],2);
    Screen('Flip',win);
    KbStrokeWait();

    % ask if the participant has questions
    questions = 'Do you have any questions?\n(Press any key to proceed)';
    DrawFormattedText(win, questions,'center', 'center', text_colour,[],[],[],2);
    Screen('Flip',win);
    KbStrokeWait();

    % END OF INSTRUCTIONS
    
    for block_counter = 1:number_of_blocks   % the loop to get through the blocks

        % Initialize stuff
        trial_counter = 1;
        
        % Start of Block message
        Screen(win,'TextFont',normal_font);
        Screen(win,'TextSize',normal_font_size*2);
        DrawFormattedText(win, ['Beginning Block  ' num2str(block_counter)],'center', 'center', text_colour,[],[],[],2);
        Screen('Flip',win);

        WaitSecs(1);

        % GO message
        Screen(win,'TextFont',normal_font);
        Screen(win,'TextSize',normal_font_size*2);
        DrawFormattedText(win, 'GO','center', 'center', text_colour,[],[],[],2);
        Screen('Flip',win);

        WaitSecs(1);

        odds = 0;   % counter for the number of odds seen
        
        % Trial loop

        while trial_counter <= trials_per_block

            % Decide if this is a oddball trial
            if (rand() < oddball_probability)
                odd_trial  = 0;
                odds = odds + 1;
                circle_colour = oddball_colour;
            else
                odd_trial = 1;
                circle_colour = control_colour;
            end
            
            % Draw fixations
            Screen('TextSize',win,fixation_size);
            DrawFormattedText(win, '+','center', 'center', fixation_colour,[],[],[],2);
            Screen('Flip',win);

            % Wait for some amount of time - note the total delay is this
            % amount (here, 100 to 300 ms + the 200 ms baseline
            jitter_amount = (rand() - 0.5) / 5;
            WaitSecs(0.2 + jitter_amount);
            
            % Clear the MUSE OSC buffer by calling osc_recv
            osc_recv(osc_stream,osc_timeout);

            % Record 200 ms of baseline
            num_samples = length_of_baseline * sample_rate;
            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;
                            % Stop looping through this message if all
                            % samples received
                            if counter > num_samples
                                break;
                            end
                        end
                    end
                end
            end

            % Draw the stimlus
            Screen('FillOval', win , circle_colour, [xmid-stim_size ymid-stim_size xmid+stim_size ymid+stim_size]); % Left
            Screen('Flip',win);

            % record EEG data for the stimulus
            num_samples = length_of_erp * sample_rate;
            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;
                            % Stop looping through this message if all
                            % samples received
                            if counter > num_samples
                                break;
                            end
                        end
                    end
                end
            end

            % Save EEG into two matrices - one for each condition
            if odd_trial == 0
                num_odd = num_odd + 1;
                this_eeg = [temp_baseline temp_eeg]; % Construct epoch
                odd_eeg = [odd_eeg this_eeg]; % Add this epoch to the oddball matrix
            else
                num_dist = num_dist + 1;
                this_eeg = [temp_baseline temp_eeg]; % Construct epoch
                dist_eeg = [dist_eeg this_eeg]; % Add this epoch to the distractor matrix
            end
            EEG{1} = odd_eeg;
            EEG{2} = dist_eeg;
            ALLEEG = [ALLEEG this_eeg]; % Store this trial's EEG
            markers = [markers odd_trial]; % Store this trial's marker
            
            % Append this trial
            data = [data; actual_trial_counter block_counter trial_counter odd_trial];

            % increment counters
            actual_trial_counter = actual_trial_counter + 1;
            trial_counter = trial_counter + 1;

            % check to see if exit has been called
            [~, ~, keyCode] = KbCheck();
            if keyCode(ExitKey)
                break;
            end
        end

        % End of block - save EEG and behavioural data
        save(datafile_name, 'EEG', 'ALLEEG', 'markers');
        save(behavioural_file_name,'data');
       
        % Display total
        Screen(win,'TextSize',normal_font_size);
        DrawFormattedText(win, ['The number of blue circles in this block was\n' num2str(odds) '\n(Press any key to proceed)'],'center', 'center', text_colour,[],[],[],2);
        Screen('Flip',win);
        
        % check to see if exit has been called
        [~, ~, keyCode] = KbCheck();
        if keyCode(ExitKey)
            break;
        end
        
        WaitSecs(1);

        % Wait until a key is pressed
        KbStrokeWait();
        
    end

catch Me
    % Close the Psychtoolbox window, bring back the cursor and keyboard,
    % free up the osc server
    Screen('CloseAll');
    ListenChar(0);
    ShowCursor();
    osc_free_server(osc_stream);
    rethrow(Me);
end

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

% Close the OSC stuff if were using the Muse
osc_free_server(osc_stream);