% This script creates figures roughly displaying the response patterns of
% the adaptation (tae) task and choice repetition (sd) task.
% It inputs the cleaned compiled data (.csv) and outputs and saves two 
% multi-panel figures, one for each task.
% It uses the Palames 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 analysis compiled data';

%% 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)];

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

for iGroup = 1:2
    
    % Calculate p_right values

    for iStimLevel = 1:length(stim_levels_tae)
        
        % Select data for this group and stimulus level
        data_select = tae_data(tae_data.group == (iGroup-2)*-1 & tae_data.stimulus == stim_levels_tae(iStimLevel),:); 
        % note: this way, iGroup 1 = ASD and iGroup 2 = control

        % all data     
        numPos(iStimLevel) = sum(data_select.decision);
        OutOfNum(iStimLevel) = height(data_select);
        p{iGroup}(iStimLevel) = numPos(iStimLevel)/OutOfNum(iStimLevel);

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

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

    end

    % 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{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', numPos, OutOfNum, searchGrid,paramsFree,PF);
    [pf_ccw{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', numPos_bias_ccw, OutOfNum_bias_ccw, searchGrid,paramsFree,PF);
    [pf_cw{iGroup}, ~, ~] = ...
        PAL_PFML_Fit(stim_levels_tae', numPos_bias_cw, OutOfNum_bias_cw, searchGrid,paramsFree,PF);
       
    % Extract and calculate response characteristics
    slope_size(iGroup) = pf_all{iGroup}(2); % save the slope (of all responses, not split)
    tae_size(iGroup) = pf_cw{iGroup}(1) - pf_ccw{iGroup}(1); % estimate the size of tae by looking at the PSE diff of split responses
    bias_ccw_cw(iGroup) = mean(p_bias_ccw{iGroup} - p_bias_cw{iGroup}, 2); % compute bias over data, without psychometric fit

end

%% TAE Plots

figure; 
set(gcf, 'Position',  [50, 200, 1500, 400])
subplot(1,3,1); hold on;
StimLevelsFineGrain = [min(stim_levels_tae):max(stim_levels_tae)./1000:max(stim_levels_tae)];
Proportion_Model_all_ASD = PF(pf_all{1},StimLevelsFineGrain);
Proportion_Model_all_TD = PF(pf_all{2},StimLevelsFineGrain);
p1 = plot(StimLevelsFineGrain,Proportion_Model_all_ASD,'-', 'Color', [0.9290, 0.6940, 0.1250]);
    plot(stim_levels_tae', p{1}, '.', 'Color', [0.9290, 0.6940, 0.1250]);
p2 = plot(StimLevelsFineGrain,Proportion_Model_all_TD,'-', 'Color', [0 0 0]);
    plot(stim_levels_tae', p{2}, '.', 'Color', [0 0 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', 'TD', 'Location', 'northwest')
legend boxoff
title('Task 1 performance') 

for iGroup = 1:2   
    subplot(1,3,iGroup+1); hold on;
    StimLevelsFineGrain     = [min(stim_levels_tae):max(stim_levels_tae)./1000:max(stim_levels_tae)];
    Proportion_Model_ccw{iGroup} =  PF(pf_ccw{iGroup},StimLevelsFineGrain);
    Proportion_Model_cw{iGroup}  =  PF(pf_cw{iGroup},StimLevelsFineGrain);
    g1 = plot(StimLevelsFineGrain,Proportion_Model_ccw{iGroup},'-r');
        plot(stim_levels_tae', p_bias_ccw{iGroup}, '.r');
    g2 = plot(StimLevelsFineGrain,Proportion_Model_cw{iGroup},'-b');
        plot(stim_levels_tae', p_bias_cw{iGroup}, '.b');
    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([g1, g2], 'ccw inducer','cw inducer', 'Location', 'northwest')
    legend boxoff
    if iGroup == 1
        group_name = 'ASD';
    else
        group_name = 'TD';
    end
    % title(['TAE ' group_name]);
    title(group_name);
end
%saveas(gcf, strcat(graph_path, '/TAE task performance plus fits ASD vs TD.png')) 

%% 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_select = sd_data(sd_data.group == (iGroup-2)*-1,:); 
    
    % Split data into previous decision = left and previous decision = right
    data_select_0 = data_select(data_select.prevDecision == 0,:);
    data_select_1 = data_select(data_select.prevDecision == 1,:);
       
    % Choice repetition probability
    p_repeat_overall{iGroup} = mean( ...
        [sum(data_select_0.prevDecision == data_select_0.currDecision) / height(data_select_0), ...
        sum(data_select_1.prevDecision == data_select_1.currDecision) / height(data_select_1)]);
    %
    %
    %
    
    % Sanity checks (ignoring serial dependence)    
    for iStimLevel = 1:length(stim_levels_sd)
        
        % Select data for this group and stimulus level
        data_select = sd_data(sd_data.group == (iGroup-2)*-1 & sd_data.currStimDir == sign(stim_levels_sd(iStimLevel)),:); 
        % note: this way, iGroup 1 = ASD and iGroup 2 = control
  
        % Proportion "right" responses
        p_right{iGroup}(:,iStimLevel) = sum(data_select.currDecision) / height(data_select);
        
        % Proportion "right" responses split by previous response
        p_right_prev_left{iGroup}(:,iStimLevel) = sum(data_select.currDecision == 1 & data_select.prevDecision == 0) / height(data_select(data_select.prevDecision == 0,:));
        p_right_prev_right{iGroup}(:,iStimLevel) = sum(data_select.currDecision == 1 & data_select.prevDecision == 1) / height(data_select(data_select.prevDecision == 1,:));
              
    end
        
end

%% SD Plots

% Task performance and split by previous response
figure; hold on;
subplot(1,3,1); hold on;
p1 = plot(stim_levels_sd', p_right{1},'.-','Color',[0.9290, 0.6940, 0.1250]);
p2 = plot(stim_levels_sd', p_right{2},'.-','Color',[0 0 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', 'TD', 'Location', 'northwest')
legend boxoff
title('Task 2 performance') 
for iGroup = 1:2 
    subplot(1,3,iGroup+1); hold on; 
    plot(stim_levels_sd', p_right_prev_left{iGroup},'.-r');
    plot(stim_levels_sd', p_right_prev_right{iGroup},'.-b');
    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('Previous response ccw','Previous response cw', 'Location', 'northwest');
    legend boxoff;
    if iGroup == 1
        group_name = 'ASD';
    else
        group_name = 'TD';
    end
    title(['SD ' group_name]);
end
set(gcf,'Position',[100 100 1200 400])
%saveas(gcf, strcat(graph_path, '/SD task performance plus responses ASD vs TD.png'))

%{
% p_repeat overall
figure; hold on;
for i=1:2
    h = bar(i, p_repeat_overall{i});
    if i == 2
        col = [0.5843 0.8157 0.9882];
        set(h, 'FaceColor', col)
    end
    clear col
end
ylim([0.45 0.55])
xlim([.3 2.7])
xticks([1 2])
set(gca,'xticklabel', {'ASD', 'TD'})
xlabel('Group')
ylabel('p(repeat)')
title('All trials');
saveas(gcf, strcat(graph_path, '/SD p_repeat ASD vs TD.png'))
%}