% This script creates figures 2 and 3 of the publication.
% It inputs the cleaned compiled data (.csv).
% It runs model-free and psychometric fitting analysis over the adaptation 
% (tae) task and  model-free analysis over the choice repetition (sd) task.
% It outputs and saves two multi-panel figures, one for each task.
% It uses the Palamedes Toolbox.
% (c) Ella Bosch, 2021

clear all; close all;
addpath ../Toolboxes/Palamedes

%% Define participant, task, and analysis specifics

% This analysis uses data with crit 0:
% 0 = exclude poor performers on SD task and people who do not meet inclusion criteria

% where do you want to save output?
path = 'P:\3018029.04\Analysis Scripts\Export';
graph_path = 'P:/3018029.04/Analysis Scripts/Graphs for manuscript';

%% TAE ANALYSIS
 
% load data
tae_data = readtable('P:\3018029.04\Data (drive)\data_tae_clean.csv');

% add column for stimulus
stimulus = tae_data.stimulus_direction .* tae_data.stimulus_evidence;
tae_data = [tae_data table(stimulus)];
clear stimulus

% read stimulus levels
stim_levels_tae = unique(tae_data.stimulus);

for iGroup = 1:2
    
    % Select data for this group
    data_group = tae_data(tae_data.group == (iGroup-2)*-1,:);
    % note: this way, iGroup 1 = ASD and iGroup 2 = control
    
    % Make a list of the number of subjects in the group
    subjects{iGroup} = unique(data_group.subject);
    
    for iSubject = 1:length(subjects{iGroup})
        
        % Select data for this subject
        data_select = tae_data(tae_data.subject == subjects{iGroup}(iSubject),:);
        
        % Calculate bias
        tae_bias{iGroup}(iSubject,:) = ...
            (height(data_select(data_select.decision == 1 & data_select.adaptor_direction == -1,:)) ...
                / height(data_select(data_select.adaptor_direction == -1,:))) ...
            - (height(data_select(data_select.decision == 1 & data_select.adaptor_direction == 1,:)) ...
                / height(data_select(data_select.adaptor_direction == 1,:)));
               
        for iStimLevel = 1:length(stim_levels_tae)
            
            % Select data for this subject and stimulus level
            data_select = tae_data(tae_data.subject == subjects{iGroup}(iSubject) & tae_data.stimulus == stim_levels_tae(iStimLevel),:); 

            % Calculate proportion of clockwise responses
            numPos(iSubject,iStimLevel) = sum(data_select.decision);
            OutOfNum(iSubject,iStimLevel) = height(data_select);
            p{iGroup}(iSubject, iStimLevel) = numPos(iSubject,iStimLevel)/OutOfNum(iSubject,iStimLevel);

            numPos_bias_ccw(iSubject,iStimLevel) = sum(data_select.decision(data_select.adaptor_direction == -1,:));
            OutOfNum_bias_ccw(iSubject,iStimLevel) = height(data_select(data_select.adaptor_direction == -1,:));
            p_bias_ccw{iGroup}(iSubject, iStimLevel) = numPos_bias_ccw(iSubject,iStimLevel)/OutOfNum_bias_ccw(iSubject,iStimLevel);

            numPos_bias_cw(iSubject,iStimLevel) = sum(data_select.decision(data_select.adaptor_direction == 1,:));
            OutOfNum_bias_cw(iSubject,iStimLevel) = height(data_select(data_select.adaptor_direction == 1,:));
            p_bias_cw{iGroup}(iSubject, iStimLevel) = numPos_bias_cw(iSubject,iStimLevel)/OutOfNum_bias_cw(iSubject,iStimLevel);

        end  
        
    end
            
    % Calculate means and SEs
    p_mean{iGroup}              = mean(p{iGroup});
        p_se{iGroup}            = std(p{iGroup})/sqrt(length(subjects{iGroup}));
    p_bias_ccw_mean{iGroup}     = mean(p_bias_ccw{iGroup});
        p_bias_ccw_se{iGroup}   = std(p_bias_ccw{iGroup})/sqrt(length(subjects{iGroup}));
    p_bias_cw_mean{iGroup}      = mean(p_bias_cw{iGroup});
        p_bias_cw_se{iGroup}    = std(p_bias_cw{iGroup})/sqrt(length(subjects{iGroup}));
        
    % Fit psychometric function
    PF = @PAL_CumulativeNormal;

    % determine free and fixed parameters, where 1 = free parameter and 0 = fixed parameter
    paramsFree = [1 1 0 0]; % correspond to threshold, slope, guess rate, and lapse rate 

    % parameter grid defining parameter space through which to perform a
    % brute-force search for values to be used as initial guesses in iterative
    % parameter search.
    searchGrid.alpha = -10:.01:10;
    searchGrid.beta = logspace(0,1,500);
    searchGrid.gamma = 0.01;  % scalar, used if fixed parameter
    searchGrid.lambda = 0.01;  % scalar, used if fixed parameter
    %searchGrid.gamma = 0:.01:.15; % vector, used if free parameter
    %searchGrid.lambda = 0:.01:.15; % vector, used if free parameter

    % Perform fit
    [pf_all_group{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', sum(numPos), sum(OutOfNum),searchGrid,paramsFree,PF);
    [pf_ccw_group{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', sum(numPos_bias_ccw), sum(OutOfNum_bias_ccw),searchGrid,paramsFree,PF);
    [pf_cw_group{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', sum(numPos_bias_cw), sum(OutOfNum_bias_cw),searchGrid,paramsFree,PF);
  
end

figure; 
StimLevelsFineGrain     = [min(stim_levels_tae):max(stim_levels_tae)./1000:max(stim_levels_tae)];

subplot(2,2,1); hold on;
Proportion_Model_ccw{1} =  PF(pf_ccw_group{1},StimLevelsFineGrain);
Proportion_Model_cw{1}  =  PF(pf_cw_group{1},StimLevelsFineGrain);
g1 = plot(StimLevelsFineGrain, Proportion_Model_ccw{1}, '-', 'Color', [0, 0.2, 0.8]);
    errorbar(stim_levels_tae', p_bias_ccw_mean{1}, p_bias_ccw_se{1}, '.', 'Color', [0, 0.2, 0.8]);
g2 = plot(StimLevelsFineGrain, Proportion_Model_cw{1}, '-', 'Color', [0.4, 0.7, 1]);
    errorbar(stim_levels_tae', p_bias_cw_mean{1}, p_bias_cw_se{1}, '.', 'Color', [0.4, 0.7, 1]);
legend([g1, g2], 'ccw inducer','cw inducer', 'Location', 'northwest')
legend boxoff
    
subplot(2,2,2); hold on;
Proportion_Model_ccw{2} =  PF(pf_ccw_group{2},StimLevelsFineGrain);
Proportion_Model_cw{2}  =  PF(pf_cw_group{2},StimLevelsFineGrain);
g1 = plot(StimLevelsFineGrain, Proportion_Model_ccw{2}, '-', 'Color', [0.8, 0.5, 0]);
    errorbar(stim_levels_tae', p_bias_ccw_mean{2}, p_bias_ccw_se{2}, '.', 'Color', [0.8, 0.5, 0]);
g2 = plot(StimLevelsFineGrain, Proportion_Model_cw{2}, '-', 'Color', [1, 0.75, 0.4]);
    errorbar(stim_levels_tae', p_bias_cw_mean{2}, p_bias_cw_se{2}, '.', 'Color', [1, 0.75, 0.4]);
legend([g1, g2], 'ccw inducer','cw inducer', 'Location', 'northwest')
legend boxoff

for iGroup = 1:2   
    subplot(2,2,iGroup); hold on;
    xlabel('Stimulus orientation from vertical (deg)')
    ylabel('p(clockwise response)')
    xlim([-14 14]);
    ylim([0 1]);
    set(gca,'XTick',stim_levels_tae)
    set(gca,'XTickLabel',{'-12';'-6';'0';'6';'12'});
    if iGroup == 1
        group_name = 'ASD group';
    else
        group_name = 'TD group';
    end
    % title(['TAE ' group_name]);
    title(group_name);
end

subplot(2,2,3); hold on;
StimLevelsFineGrain = [min(stim_levels_tae):max(stim_levels_tae)./1000:max(stim_levels_tae)];
Proportion_Model_all_ASD = PF(pf_all_group{1},StimLevelsFineGrain);
Proportion_Model_all_TD = PF(pf_all_group{2},StimLevelsFineGrain);
p1 = plot(StimLevelsFineGrain,Proportion_Model_all_ASD,'-', 'Color', [0, 0.5, 1]); 
    errorbar(stim_levels_tae', p_mean{1}, p_se{1}, '.', 'Color', [0, 0.5, 1]);
p2 = plot(StimLevelsFineGrain,Proportion_Model_all_TD,'-', 'Color', [0.9, 0.6, 0]); 
    errorbar(stim_levels_tae', p_mean{2}, p_se{2}, '.', 'Color', [0.9, 0.6, 0]);
xlabel('Stimulus orientation from vertical (deg)')
ylabel('p(clockwise response)')
xlim([-14 14]);
ylim([0 1]);
set(gca,'XTick',stim_levels_tae)
set(gca,'XTickLabel',{'-12';'-6';'0';'6';'12'});
legend([p1, p2], 'ASD group', 'TD group', 'Location', 'northwest')
legend boxoff
title('Overall task performance') 

subplot(2,2,4); hold on;
bias.ASD    = tae_bias{1};
bias.TD     = tae_bias{2};
v = violinplot(bias);
v(1).ViolinColor = [0, 0.5, 1];
v(2).ViolinColor = [0.9, 0.6, 0];
ylim([-0.1 0.65])
ylabel('\Delta proportion cw responses')
title('Size of adaptation bias')

set(gcf, 'Position',  [100, 100, 800, 800])
saveas(gcf, strcat(graph_path, '/TAE task performance plus fits ASD vs TD.png')) 
saveas(gcf, strcat(graph_path, '/TAE task performance plus fits ASD vs TD.pdf')) 

v(1).ViolinAlpha = 1;
v(2).ViolinAlpha = 1;
saveas(gcf, strcat(graph_path, '/TAE task performance plus fits ASD vs TD.eps')) 

% Clear some variables
clear data_group subjects

%% SD ANALYSIS

% load data
sd_data = readtable('P:\3018029.04\Data (drive)\data_sd_clean.csv');

% add column for current stimulus
currStimulus = sd_data.currStimDir .* sd_data.currStimEvidence;
sd_data = [sd_data table(currStimulus)];

% read stimulus levels
stim_levels_sd = unique(sd_data.currStimulus);

for iGroup = 1:2
    
    % Select data for this group
    data_group = sd_data(sd_data.group == (iGroup-2)*-1,:);
    % note: this way, iGroup 1 = ASD and iGroup 2 = control
    
    % Make a list of the number of subjects in the group
    subjects{iGroup} = unique(data_group.subject);
    
    for iSubject = 1:length(subjects{iGroup})
        
            % Select subject data
            data_select = sd_data(sd_data.subject == subjects{iGroup}(iSubject),:);
    
            % Calculate p(repeat)
            p_repeat{iGroup}(iSubject,:) = mean( ...
                [ height(data_select(data_select.currDecision == data_select.prevDecision & data_select.prevDecision == 0,:)) ...
                    / height(data_select(data_select.prevDecision == 0,:)), ...
                height(data_select(data_select.currDecision == data_select.prevDecision & data_select.prevDecision == 1,:)) ...
                    / height(data_select(data_select.prevDecision == 1,:)) ]);

        for iStimLevel = 1:length(stim_levels_sd)
        
            % Select data for this subject and stimulus level
            data_select = sd_data(sd_data.subject == subjects{iGroup}(iSubject) & sd_data.currStimDir == sign(stim_levels_sd(iStimLevel)),:); 
            % note: this way, iGroup 1 = ASD and iGroup 2 = control

            % Proportion "right" responses
            p_right{iGroup}(iSubject,iStimLevel) = sum(data_select.currDecision) / height(data_select);

            % Proportion "right" responses split by previous response
            p_right_prev_left{iGroup}(iSubject,iStimLevel) = sum(data_select.currDecision == 1 & data_select.prevDecision == 0) / height(data_select(data_select.prevDecision == 0,:));
            p_right_prev_right{iGroup}(iSubject,iStimLevel) = sum(data_select.currDecision == 1 & data_select.prevDecision == 1) / height(data_select(data_select.prevDecision == 1,:));

        end
        
    end
    
    % Calculate means and SEs   
    p_right_mean{iGroup}                    = mean(p_right{iGroup});
        p_right_se{iGroup}                  = std(p_right{iGroup})/sqrt(length(subjects{iGroup}));     
    p_right_prev_left_mean{iGroup}          = mean(p_right_prev_left{iGroup});
        p_right_prev_left_se{iGroup}        = std(p_right_prev_left{iGroup})/sqrt(length(subjects{iGroup}));
    p_right_prev_right_mean{iGroup}         = mean(p_right_prev_right{iGroup});
        p_right_prev_right_se{iGroup}       = std(p_right_prev_right{iGroup})/sqrt(length(subjects{iGroup}));
    
end

%% SD Plots

% Task performance and split by previous response
figure; hold on;

subplot(2,2,1); hold on; 
errorbar(stim_levels_sd', p_right_prev_left_mean{1}, p_right_prev_right_se{1}, '.-', 'Color', [0, 0.2, 0.8]);
errorbar(stim_levels_sd', p_right_prev_right_mean{1}, p_right_prev_right_se{1}, '.-', 'Color', [0.4, 0.7, 1]);
legend('Previous response ccw','Previous response cw', 'Location', 'northwest');
legend boxoff;

subplot(2,2,2); hold on; 
errorbar(stim_levels_sd', p_right_prev_left_mean{2}, p_right_prev_right_se{2}, '.-', 'Color', [0.8, 0.5, 0]);
errorbar(stim_levels_sd', p_right_prev_right_mean{2}, p_right_prev_right_se{2}, '.-', 'Color', [1, 0.75, 0.4]);
legend('Previous response ccw','Previous response cw', 'Location', 'northwest');
legend boxoff;
    
for iGroup = 1:2 
    subplot(2,2,iGroup); hold on; 
    xlabel('Stimulus orientation from vertical (deg)')
    ylabel('p(clockwise response)')
    xlim([-6 6]);
    ylim([0 1]);
    set(gca,'XTick',stim_levels_sd')
    set(gca,'XTickLabel',{'-5';'0';'5'});
    if iGroup == 1
        group_name = 'ASD group';
    else
        group_name = 'TD group';
    end
    title(group_name);
end

subplot(2,2,3); hold on;
p1 = errorbar(stim_levels_sd', p_right_mean{1}, p_right_se{1}, '.-','Color', [0, 0.5, 1]);
p2 = errorbar(stim_levels_sd', p_right_mean{2}, p_right_se{2}, '.-','Color', [0.9, 0.6, 0]);
xlabel('Stimulus orientation from vertical (deg)')
ylabel('p(clockwise response)')
xlim([-6 6]);
ylim([0 1]);
set(gca,'XTick',stim_levels_sd');
set(gca,'XTickLabel',{'-5';'0';'5'});
legend([p1, p2], 'ASD group', 'TD group', 'Location', 'northwest')
legend boxoff
title('Overall task performance') 

subplot(2,2,4); hold on;
bias.ASD    = p_repeat{1};
bias.TD     = p_repeat{2};
violinAlpha = 0;
v = violinplot(bias);
v(1).ViolinColor = [0, 0.5, 1];
v(2).ViolinColor = [0.9, 0.6, 0];
ylim([0.4 0.72])
ylabel('p(repeat)')
title('Size of choice repetition bias')

set(gcf,'Position',[100 100 800 800])
saveas(gcf, strcat(graph_path, '/SD task performance plus responses ASD vs TD.png'))
saveas(gcf, strcat(graph_path, '/SD task performance plus responses ASD vs TD.pdf'))

v(1).ViolinAlpha = 1;
v(2).ViolinAlpha = 1;
saveas(gcf, strcat(graph_path, '/SD task performance plus responses ASD vs TD.eps'))
