% =========================================================================
% EEG project DCCN 2015.
% Data preparation for mixed-model analysis of motivational Go/NoGo 
% learning task.
% By J.C. Swart, 2014.
% =========================================================================
function EEGpav_prepR
dbstop if error
% 1. SETUP.
% -------------------------------------------------------------------------
[data,par,dirs] = EEGpav_setup;

% 2. RETRIEVE DATA FROM FILES.
% ---------------------------------------------------------------------
if ~exist(dirs.singletrial,'file')
    EEGpav_getdata(data,par,dirs);
end

% 3. PREPARE CSV FILE FOR MIXED-MODEL ANALYSIS IN R.
% ---------------------------------------------------------------------
if ~exist(dirs.csv,'file')
    EEGpav_prepmixedmodel4R(dirs);
end

% 4. PREPARE DATA FOR STAN IN R.
% ---------------------------------------------------------------------
if ~exist(dirs.stan,'file')
    EEGpav_prepStan4R(dirs)
end

% -------------------------------------------------------------------------
disp('End of script.')
% -------------------------------------------------------------------------
end % end EEGpav_prepsingletrialanalysis.
% -------------------------------------------------------------------------
function [data,par,dirs] = EEGpav_setup
% 1. SETUP.
% -------------------------------------------------------------------------

% setup directories.
if isunix
    keyboard;
elseif ispc
    base = 'M:\';
end
if ~exist('matlabScripts','dir') % add matlab script for writing to csv.
    addpath(genpath(fullfile(base,'matlabScripts')));
end
dirs.root       = fullfile(base,'projects','EEG');
dirs.log        = fullfile(dirs.root,'Log','behaviour');
dirs.singletrial= fullfile(dirs.root,'results','EEGsingletrial.mat');
dirs.csv        = fullfile(dirs.root,'results','EEG4mixedmodelR.csv');
dirs.stan       = fullfile(dirs.root,'results','EEG4stan.mat');
cd(dirs.log);
dirs.sub        = ls('3017033.03_jesmaa_*_001');
clear base

% task parameters.
par.nSub        = size(dirs.sub,1);
par.nPart       = 2;
par.nTrial      = 320;
par.nCue        = 8;

% initialize.
par.subjects    = zeros(par.nSub,1);
data.acc        = NaN(par.nSub,par.nTrial*par.nPart);
data.go         = NaN(size(data.acc));
data.outcome    = NaN(size(data.acc));
data.RT         = NaN(size(data.acc));
data.stim       = NaN(size(data.acc));
data.resp       = NaN(size(data.acc));
data.reqResp    = NaN(size(data.acc));

end
% -------------------------------------------------------------------------
function EEGpav_getdata(data,par,dirs)
% 2. RETRIEVE DATA FROM FILES.
% -------------------------------------------------------------------------
  
% loop over files.
for iFile = 1:length(dirs.sub)
    
    % load subject data.
    sID     = dirs.sub(iFile,end-6:end-4);
    
    filename = fullfile(dirs.log,dirs.sub(iFile,:),[dirs.sub(iFile,:) '_results.mat']);
    
    if ~exist(filename,'file')
        continue
    end
    
    load(filename);
    par.subjects(iFile)    = str2double(sID);
    
    for iPart = 1:par.nPart
        % store subject info.
        idx = 1:par.nTrial; if iPart == 2; idx = idx + par.nTrial; end
        data.go(iFile,idx)       = results.learn{iPart}.go;
        data.acc(iFile,idx)      = results.learn{iPart}.acc;
        data.RT(iFile,idx)       = results.learn{iPart}.RT;
        data.outcome(iFile,idx)  = results.learn{iPart}.outcome;
        data.stim(iFile,idx)     = prep.seq.learn.stim{iPart};
        if iPart==2; data.stim(iFile,idx) = data.stim(iFile,idx) + par.nCue; end 
        data.reqResp(iFile,idx)  = prep.seq.learn.resp{iPart};
        tmp                     = results.learn{iPart}.response;
        tmp(tmp == prep.par.key.left)  = 1;
        tmp(tmp == prep.par.key.right) = 2;
        data.resp(iFile,idx)     = tmp;
        
    end % end iPart-loop.
    
end % end iFile-loop.

save(dirs.singletrial,'data','par','dirs')

end
% -------------------------------------------------------------------------
function EEGpav_prepmixedmodel4R(dirs)
% 3. PREPARE CSV FILE FOR MIXED-MODEL ANALYSIS IN R.
% ---------------------------------------------------------------------
% NOTE: Data should be sorted, so that each column contains a DV, IV or
% covariate. Each column contains for that variable the data of each trial,
% first for MPH then placebo, for subject by subject.

% load single-trial data.
load(dirs.singletrial)

% number of data points.
nRow        = numel(data.go);

% ivs.
% ---------------------------------------------------------------------
% subjects.
sID         = reshape(repmat(par.subjects',par.nTrial*par.nPart,1),nRow,1);

% cue valence, required action, and trial index.
cueID       = reshape(data.stim',nRow,1);
rewardcues  = [1 2 5 6 9 10 13 14];
valenceID   = ismember(cueID,rewardcues);
gocues      = [1:4 9:12];
actionID    = ismember(cueID,gocues);
trialID     = repmat([1:par.nTrial]',par.nSub*par.nPart,1);

% combine in matrix.
ivmatrix    = [sID valenceID actionID trialID];

% dvs: Go, accuracy, lnRT.
% ---------------------------------------------------------------------
goID        = reshape(data.go',nRow,1);
accID       = reshape(data.acc',nRow,1);
rtID        = reshape(data.RT',nRow,1);

% combine dvs in matrix.
dvmatrix    = [goID accID rtID];

% combine matrices.
datamatrix  = [ivmatrix dvmatrix];
columnlabels= {'sID,','valence,','action,','trial,','DVgo,','DVacc,','DVrt,'};

% write the data to csv.
hdrAll      = cat(2,columnlabels{:});

mycsvwrite(dirs.csv,datamatrix,hdrAll(1:end-1));

end
% -------------------------------------------------------------------------
function EEGpav_prepStan4R(dirs)
% 4. PREPARE DATA FOR STAN IN R.
% ---------------------------------------------------------------------
% ya/a/sb/rb/rew is matrix nSub x nTrial

% backup stan directory.
standir = dirs.stan;

% load single-trial data.
load(dirs.singletrial)

% number of data points.
nRow        = numel(data.go);

% ya, actual action that was performed, 1 = go, 0 = nogo.
ya              = data.resp;
ya(ya == 99)    = NaN; % wrong button.
ya(ya == 98)    = NaN; % wrong button.
ya(ya == 102)   = NaN; % wrong button.
ya(ya == 65)    = 1; % release of left button.
ya(ya == 69)    = 2; % release of right button.

% a, recode nogo as 3 for indexing.
a               = ya; 
a(ya == 0)      = 3; 
% s, stimulus ID on each trial (stimID+8 for PLA cues).
s               = data.stim;
% r, experienced outcome on each trial, {-1, 0, 1}.
r               = data.outcome;
% rew, cue-valence, 1 = win-cue, 0 = avoid-cue.
rew             = NaN(size(ya));
rew(ismember(s,[1 2 5 6 9 10 13 14]))   = 1;
rew(ismember(s,[3 4 7 8 11 12 15 16]))  = 0;
% N, number of trials = 2 x 320.
N               = size(ya,2);
% Nsub, number of subjects = 106.
Nsub            = size(ya,1);

% and save the data for rStan.
save(standir,'ya','a','s','r','rew','N','Nsub')

end
% -------------------------------------------------------------------------
