% This script is used to run a rough single-subject analysis.
% It calls on scripts that load the behavioral task files from the
% choice repetition (sd) task and adaptation (tae) task of one (or multiple)
% individual participant(s).
% It then calculates several dependent measures and plots the response
% patterns for both tasks, saving these figures. 
% (c) Ella Bosch, 2021

clear all; close all;

%% Define what participants

subjects = 45;

stim_levels = -5:5:5;

cfg = [];
cfg.subjects = subjects;
cfg.plot = false;

%% SD Analysis

% Load data
[sd_data, sd_log] = sd_load_data(cfg);

% Sanity Checks (ignoring serial dependence)
% How do the responses on the current trial relate to the shown stimuli
for iSubject = 1:length(subjects)
    
    % stimulus left
    p_right{iSubject}(1) = nansum(sd_data{iSubject}(sd_data{iSubject}(:,5) == -5,8)) / ...
                                sum(~isnan(sd_data{iSubject}(sd_data{iSubject}(:,5) == -5,8)));
                
    % stimulus zero                       
    p_right{iSubject}(2) = nansum(sd_data{iSubject}(sd_data{iSubject}(:,5) == 0,8)) / ...
                                sum(~isnan(sd_data{iSubject}(sd_data{iSubject}(:,5) == 0,8))); 
                            
    % stimulus right                       
    p_right{iSubject}(3) = nansum(sd_data{iSubject}(sd_data{iSubject}(:,5) == 5,8)) / ...
                                sum(~isnan(sd_data{iSubject}(sd_data{iSubject}(:,5) == 5,8)));                        
    
                            
    figure; bar(p_right{iSubject}); 
    ylim([0 1]); 
    xlim([0 4]); 
    set(gca,'XTick',1:3,'XTickLabel',{'-5','0','5'});
    xlabel('Stimulus orientation')
    ylabel('P(right)')
    
    title(['Subject ' num2str(subjects(iSubject))]);
                            
end


% Sort data for serial dependence analyses (for all subjects)
cfg = [];
cfg.subjects    = subjects;
cfg.data        = sd_data;
cfg.plot        = true;

sorted_sd_data  = sd_sort_data(cfg);

% Compute probability of repeating the decision of the previous trial
for iSubject = 1:length(subjects)
   
    % probability to repeat the decision (regardless of stimulus)
    p_repeat_decision(iSubject) = sum(sorted_sd_data{iSubject}(:,2) == sorted_sd_data{iSubject}(:,4)) / size(sorted_sd_data{iSubject},1);
    
    % probability to repeat the decision when current stimulus is ambiguous
    p_repeat_decision_curr_ambig(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,2) == sorted_sd_data{iSubject}(:,4) & ...
        sorted_sd_data{iSubject}(:,3) == 0)...
        / sum(sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is left
    p_repeat_decision_curr_ambig_prev_L(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == sorted_sd_data{iSubject}(:,2)) ...
        /  sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is left, and was correctly identified
    p_repeat_decision_curr_ambig_prev_L_corr(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,2) == 0 & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == 0) ...
        /  sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,2) == 0 & ...
        sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is right
    p_repeat_decision_curr_ambig_prev_R(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == sorted_sd_data{iSubject}(:,2)) ...
        /  sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is right, and was correctly identified
    p_repeat_decision_curr_ambig_prev_R_corr(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,2) == 1 & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == 1) ...
        /  sum(...
        sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,2) == 1 & ...
        sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is left or right
    p_repeat_decision_curr_ambig_prev_umambig(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,2) == sorted_sd_data{iSubject}(:,4) & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,1) ~= 0) ...
        / sum(sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,1) ~= 0);
    
    % probability to repeat the decision when current stimulus is ambiguous
    % and previous stimulus is left or right, and was correctly identified
    p_repeat_decision_curr_ambig_prev_unambig_corr(iSubject) = (...
        sum(sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,2) == 0 & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == 0) ...
        + sum(sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,2) == 1 & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,4) == 1)) ...
        / ...
        (sum(sorted_sd_data{iSubject}(:,1) == stim_levels(1) & ...
        sorted_sd_data{iSubject}(:,2) == 0 & ...
        sorted_sd_data{iSubject}(:,3) == 0) ...
        + sum(sorted_sd_data{iSubject}(:,1) == stim_levels(3) & ...
        sorted_sd_data{iSubject}(:,2) == 1 & ...
        sorted_sd_data{iSubject}(:,3) == 0));
    
    
    % probability to repeat the decision when current and previous stimuli are ambiguous
    p_repeat_decision_curr_ambig_prev_ambig(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,2) == sorted_sd_data{iSubject}(:,4) & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,1) == 0) ...
        / sum(sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,1) == 0);
    
    summary = {'p_repeat_decision', p_repeat_decision(iSubject);
        'p_repeat_decision_curr_ambig', p_repeat_decision_curr_ambig(iSubject);
        'p_repeat_decision_curr_ambig_prev_L', p_repeat_decision_curr_ambig_prev_L(iSubject);
        'p_repeat_decision_curr_ambig_prev_L_corr', p_repeat_decision_curr_ambig_prev_L_corr(iSubject);
        'p_repeat_decision_curr_ambig_prev_R', p_repeat_decision_curr_ambig_prev_R(iSubject);
        'p_repeat_decision_curr_ambig_prev_R_corr', p_repeat_decision_curr_ambig_prev_R_corr(iSubject);
        'p_repeat_decision_curr_ambig_prev_umambig', p_repeat_decision_curr_ambig_prev_umambig(iSubject);
        'p_repeat_decision_curr_ambig_prev_unambig_corr', p_repeat_decision_curr_ambig_prev_unambig_corr(iSubject);
        'p_repeat_decision_curr_ambig_prev_ambig', p_repeat_decision_curr_ambig_prev_ambig(iSubject)};
             
   
    
    % different visualisation

    for iStimLevel = 1:length(stim_levels)

        p_all{iSubject}(iStimLevel)               = nansum(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel),4)) ...
                                                    / sum(~isnan(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel),4)));
        p_prev_left{iSubject}(iStimLevel)         = nansum(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_sd_data{iSubject}(:,2) == 0,4)) ...
                                                    / sum(~isnan(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_sd_data{iSubject}(:,2) == 0,4)));
        p_prev_right{iSubject}(iStimLevel)        = nansum(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_sd_data{iSubject}(:,2) == 1,4)) ...
                                                    / sum(~isnan(sorted_sd_data{iSubject}(sorted_sd_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_sd_data{iSubject}(:,2) == 1,4)));
    end


    %% Fit models
    fits.model_all{iSubject}   = fit(stim_levels.',p_all{iSubject}.','smoothingspline');
    fits.model_left{iSubject}  = fit(stim_levels.',p_prev_left{iSubject}.','smoothingspline');
    fits.model_right{iSubject} = fit(stim_levels.',p_prev_right{iSubject}.','smoothingspline');

    %% Write in output structure

    fits.all{iSubject}        = p_all{iSubject};
    fits.prev_left{iSubject}  = p_prev_left{iSubject};
    fits.prev_right{iSubject} = p_prev_right{iSubject};

    %% Figure

    if cfg.plot == true

        figure; hold on;
        plot(stim_levels,fits.prev_left{iSubject},'+r');
        plot(stim_levels,fits.prev_right{iSubject},'+b');
        plot(fits.model_left{iSubject},'-r');
        plot(fits.model_right{iSubject},'-b');

        xlabel('Stimulus orientation')
        ylabel('P("right" response)')

        xlim([-6 6]);
        ylim([0 1]);

        set(gca,'XTick',stim_levels)
        set(gca,'XTickLabel',{'-5';'0';'5'});

        legend('Previous response "left"','Previous response "right"');

        title(['Subject ' num2str(subjects(iSubject))]);

        % saveas(gcf,['Figures/SingleSubjectFits/sd_ss_fit_' num2str(subject)'.png'])

    end
    
     
end

% Compute probability of repeating the motor response of the previous trial
for iSubject = 1:length(subjects)
    
    % probability to repeat the decision (regardless of stimulus)
    p_repeat_motor(iSubject) = sum(sorted_sd_data{iSubject}(:,5) == sorted_sd_data{iSubject}(:,6)) / size(sorted_sd_data{iSubject},1);
    
    % probability to repeat the decision when current stimulus is ambiguous
    p_repeat_motor_curr_ambig(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,5) == sorted_sd_data{iSubject}(:,6) & ...
        sorted_sd_data{iSubject}(:,3) == 0)...
        / sum(sorted_sd_data{iSubject}(:,3) == 0);
    
    % probability to repeat the decision when current and previous stimuli are ambiguous
    p_repeat_motor_curr_ambig_prev_ambig(iSubject) = sum(...
        sorted_sd_data{iSubject}(:,5) == sorted_sd_data{iSubject}(:,6) & ...
        sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:,1) == 0) ...
        / sum(sorted_sd_data{iSubject}(:,3) == 0 & ...
        sorted_sd_data{iSubject}(:, 1) == 0);
    
end
%}

%% TAE Analysis

% Load data
[tae_data, log] = tae_load_data(cfg);


cfg = [];
cfg.subjects    = subjects;
cfg.plot = true;

% Process data
for iSubject = 1:length(subjects)
    
    cfg.subject = subjects(iSubject);
    
    cfg.data = tae_data{iSubject};
    
    %cfg.savepath = ['Matlab datafiles/single subject fits/subject_' num2str(subjects(iSubject)) '_fits.mat'];
         
    tae_ss_pf{iSubject} = tae_ss_fit_pf(cfg);

end


 
% Sort data for adaptation analyses (for all subjects)
cfg = [];
cfg.subjects    = subjects;
cfg.data        = tae_data;
cfg.plot        = true;

stim_levels = -12:6:12;

for iSubject = 1:length(subjects)

    sorted_tae_data  = tae_sort_data(cfg);
    
% different visualisation

    for iStimLevel = 1:length(stim_levels)

        p_all{iSubject}(iStimLevel)               = nansum(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel),4)) ...
                                                    / sum(~isnan(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel),4)));
        p_prev_left{iSubject}(iStimLevel)         = nansum(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_tae_data{iSubject}(:,2) == 0,4)) ...
                                                    / sum(~isnan(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_tae_data{iSubject}(:,2) == 0,4)));
        p_prev_right{iSubject}(iStimLevel)        = nansum(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_tae_data{iSubject}(:,2) == 1,4)) ...
                                                    / sum(~isnan(sorted_tae_data{iSubject}(sorted_tae_data{iSubject}(:,3) == stim_levels(iStimLevel) & sorted_tae_data{iSubject}(:,2) == 1,4)));
    end

    %% Fit models
    fits.model_all{iSubject}   = fit(stim_levels.',p_all{iSubject}.','smoothingspline');
    fits.model_left{iSubject}  = fit(stim_levels.',p_prev_left{iSubject}.','smoothingspline');
    fits.model_right{iSubject} = fit(stim_levels.',p_prev_right{iSubject}.','smoothingspline');

    %% Write in output structure

    fits.all{iSubject}        = p_all{iSubject};
    fits.prev_left{iSubject}  = p_prev_left{iSubject};
    fits.prev_right{iSubject} = p_prev_right{iSubject};

    %% Figure

    if cfg.plot == true

        figure; hold on;
        plot(stim_levels,fits.prev_left{iSubject},'+r');
        plot(stim_levels,fits.prev_right{iSubject},'+b');
        plot(fits.model_left{iSubject},'-r');
        plot(fits.model_right{iSubject},'-b');

        xlabel('Stimulus orientation')
        ylabel('P("right" response)')

        xlim([-13 13]);
        ylim([0 1]);

        set(gca,'XTick',stim_levels)
        set(gca,'XTickLabel',{'-12';'-6';'0';'6';'12'});

        legend('Previous response "left"','Previous response "right"');

        title(['Subject ' num2str(subjects(iSubject))]);

        % saveas(gcf,['Figures/SingleSubjectFits/sd_ss_fit_' num2str(subject)'.png'])

    end
    
end
    
