%% fig. S6C-F - Correlation histogram and cumulatiive density function 
% Load necessary data for fig.S6C-D
clear all
close all
% =========================================================================
%% Load data
% =========================================================================
%Prepare variables
%load('figS6CE.mat');  %For locomotion
load('figS6DF.mat'); %For pupil 

% =========================================================================
%% r value distribution - fig. S6C
% =========================================================================

fig1 = figure('Name', 'Hit_CorrelationHistogram');
fig1.Position = [680 500 375 450];
BinSize = 0.05;
Edges = -1:BinSize:1;%BinSize

for iPhase = 1:4
    %Stim------------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase);
            histogram(StimCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 2
            subplot(3, 2, iPhase-1);
            histogram(StimCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [1 0 0], 'EdgeColor', 'none');
        case 3
            subplot(3, 2, iPhase-1);
            histogram(StimCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 4
            subplot(3, 2, iPhase-2);
            histogram(StimCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [1 0 0], 'EdgeColor', 'none');        
    end
    
    xlim([-1 1]);
    xticks([-1 -0.5 0 0.5 1]);
    ylim([0 0.5]);
    yticks([0 0.2 0.4]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    box off
    hold on
    switch iPhase
        case 1
            title('Stimulus Corr.')    
        case 3
            title('REV-Stimulus Corr.')
    end
    
    %Delay-----------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase+2);
            histogram(DelayCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 2
            subplot(3, 2, iPhase-1+2);
            histogram(DelayCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [0 1 0],'EdgeColor', 'none');
        case 3
            subplot(3, 2, iPhase-1+2);
            histogram(DelayCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 4
            subplot(3, 2, iPhase-2+2);
            histogram(DelayCorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [0 1 0], 'EdgeColor', 'none');        
    end
    xlim([-1 1]);
    xticks([-1 -0.5 0 0.5 1]);
    ylim([0 0.5]);
    yticks([0 0.2 0.4]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    box off
    hold on
    switch iPhase
        case 1
            title('Delay Corr.')
        case 3
            title('REV-Delay Corr.')
    end
    
    %ITI-------------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase+4);
            histogram(ITICorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 2
            subplot(3, 2, iPhase-1+4);
            histogram(ITICorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [0 0 1], 'EdgeColor', 'none');
        case 3
            subplot(3, 2, iPhase-1+4);
            histogram(ITICorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', 'none');
        case 4
            subplot(3, 2, iPhase-2+4);
            histogram(ITICorrHit_RVal(:,iPhase), Edges, 'Normalization', 'Probability', 'FaceColor', [0 0 1], 'EdgeColor', 'none');        
    end
    xlim([-1 1]);
    xticks([-1 -0.5 0 0.5 1]);
    ylim([0 0.5]);
    yticks([0 0.2 0.4]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    xlabel('Corrlation coefficient')
    box off
    hold on
    switch iPhase
        case 1
            title('ITI Corr.')
        case 3
            title('REV-ITI Corr.')
    end
end%iPhase
suptitle('Hit trial correlation-Edgeblack = Naive')

% =================================================================
%% cdf plot - fig. S6C
% =================================================================
fig2 = figure('Name', 'HitCDF');
fig2.Position = [680 500 375 450];
for iPhase = 1:4
    %Stim------------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase);
            h1 = cdfplot(abs(StimCorrHit_RVal(:,iPhase))); 
            set(h1, 'Color', 'k', 'LineWidth', 1.0, 'LineStyle', '-');
        case 2
            sub1 = subplot(3, 2, iPhase-1); %To remove XLabel, handle is necessary
            h2 = cdfplot(abs(StimCorrHit_RVal(:,iPhase))); 
            set(h2, 'Color', [1 0 0], 'LineWidth', 1.5, 'LineStyle', '-');
            sub1.XLabel.String = [];
        case 3
            subplot(3, 2, iPhase-1);
            h1 = cdfplot(abs(StimCorrHit_RVal(:,iPhase)));
            set(h1, 'Color', 'k', 'LineWidth', 1.0,'LineStyle', '-');
        case 4
            sub2 = subplot(3, 2, iPhase-2);
            h2 = cdfplot(abs(StimCorrHit_RVal(:,iPhase))); 
            set(h2, 'Color', [1 0 0], 'LineWidth', 1.5, 'LineStyle', '-');
            sub2.XLabel.String = [];
    end
    
    xlim([0 1]);
    xticks([0 0.5 1]);
    ylim([0 1.0]);
    yticks([0 0.5 1.0]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    box off
    hold on
    
    %Due to cdfplot, use case 2 and case 4 to add the tile
    switch iPhase
        case 2
             title('Stim Corr')
        case 4
             title('Stim Corr')
    end
    
    %Delay-----------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase+2);
            h1 = cdfplot(abs(DelayCorrHit_RVal(:,iPhase))); 
            set(h1, 'Color', 'k', 'LineWidth', 1.0, 'LineStyle', '-');
        case 2
            sub1 = subplot(3, 2, iPhase-1+2);
            h2 = cdfplot(abs(DelayCorrHit_RVal(:,iPhase)));
            set(h2, 'Color', [0 1 0], 'LineWidth', 1.5, 'LineStyle', '-');
            sub1.XLabel.String = [];
        case 3
            subplot(3, 2, iPhase-1+2);
            h1 = cdfplot(abs(DelayCorrHit_RVal(:,iPhase))); 
            set(h1, 'Color', 'k', 'LineWidth', 1.0,'LineStyle', '-');
        case 4
            sub2 = subplot(3, 2, iPhase-2+2);
            h2 = cdfplot(abs(DelayCorrHit_RVal(:,iPhase))); 
            set(h2, 'Color', [0 1 0], 'LineWidth', 1.5, 'LineStyle', '-');
            sub2.XLabel.String = [];
    end
    
    xlim([0 1]);
    xticks([0 0.5 1]);
    ylim([0 1.0]);
    yticks([0 0.5 1.0]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    box off
    hold on
    
    %Due to cdfplot, use case 2 and case 4 to add the tile
    switch iPhase
        case 2
            title('Delay Corr')
        case 4
            title('REV-Delay Corr')
    end
    
    %ITI-------------------------------------------------------------------
    switch iPhase
        case 1
            subplot(3, 2, iPhase+4);
            h1 = cdfplot(abs(ITICorrHit_RVal(:,iPhase))); 
            set(h1, 'Color', 'k', 'LineWidth', 1.0, 'LineStyle', '-');
        case 2
            sub1 = subplot(3, 2, iPhase-1+4);
            h2 = cdfplot(abs(ITICorrHit_RVal(:,iPhase))); 
            set(h2, 'Color', [0 0 1], 'LineWidth', 1.5, 'LineStyle', '-');
            sub1.XLabel.String = 'Absolute r value';
        case 3
            subplot(3, 2, iPhase-1+4);
            h1 = cdfplot(abs(ITICorrHit_RVal(:,iPhase))); 
            set(h1, 'Color', 'k', 'LineWidth', 1.0,'LineStyle', '-');
        case 4
            sub2 = subplot(3, 2, iPhase-2+4);
            h2 = cdfplot(abs(ITICorrHit_RVal(:,iPhase))); 
            set(h2, 'Color', [0 0 1], 'LineWidth', 1.5, 'LineStyle', '-');
            sub2.XLabel.String = 'Absolute r value';
    end
    
    xlim([0 1]);
    xticks([0 0.5 1]);
    ylim([0 1.0]);
    yticks([0 0.5 1.0]);
    ytix = get(gca, 'YTick');
    set(gca, 'YTick',ytix, 'YTickLabel',ytix*100)
    ylabel('Proportion of cell (%)')
    box off
    hold on

    %Due to cdfplot, use case 2 and case 4 to add the tile
    switch iPhase
        case 2
            title('ITI Corr')
        case 4
            title('Rev-ITI Corr')
    end
end%iPhase
suptitle('Hit trial correlation CDF -Edgeblack = Naive')

% =================================================================
%% Bar plot - fig. S6E, Mean+/-SEM for delta abs-Expert - abs-Naive
% =================================================================
NofCell = 210;
catNames = {'Stimulus', 'Delay', 'ITI'};
DeltaStim  = [];
DeltaDelay = [];
DeltaITI   = []; 
MeanDeltaStim = [];
SEM4Stim = [];

%Absolute R value in Expert - abs R in Naive for all matched cells
for iPhase = [1 3] %1: Initial, 3: Reversal
    
    DeltaStim     = abs(StimCorrHit_RVal(:, iPhase+1))-abs(StimCorrHit_RVal(:, iPhase));
    MeanDeltaStim = mean(DeltaStim, 1);
    SEM4Stim      = std(DeltaStim, 1)/sqrt(NofCell);
    
    DeltaDelay     = abs(DelayCorrHit_RVal(:, iPhase+1))-abs(DelayCorrHit_RVal(:, iPhase));
    MeanDeltaDelay = mean(DeltaDelay, 1);
    SEM4Delay      = std(DeltaDelay, 1)/sqrt(NofCell);
    
    DeltaITI     = abs(ITICorrHit_RVal(:, iPhase+1))-abs(ITICorrHit_RVal(:, iPhase));
    MeanDeltaITI = mean(DeltaITI, 1);
    SEM4ITI      = std(DeltaITI, 1)/sqrt(NofCell);
    
    
    %Bar chart-----------------------------------------------------------------
    switch iPhase
        case 1
            FigTitle   = 'Initial-delta bar chart';
            Figtitle_Inside = 'Initial - Delta r value';
            
        case 3
            FigTitle   = 'Reversal-delta bar chart';
            Figtitle_Inside = 'Reversal - Delta r value';
    end
    
    fig = figure('Name', FigTitle);
    fig.Position = [410 278 300 125];
    
    Mean4Bar = [MeanDeltaStim MeanDeltaDelay MeanDeltaITI];
    SEM4Bar  = [SEM4Stim SEM4Delay SEM4ITI];
    
    bar(1:size(catNames,2), Mean4Bar, 'FaceColor', 'none', 'EdgeColor', 'k');
    hold on;
    er = errorbar(1:size(catNames,2), Mean4Bar, SEM4Bar);
    er.Color = [0 0 0];
    er.LineStyle = 'none';
    er.MarkerEdgeColor = 'none';
    
    xticks([1 2 3 4]);
    xticklabels(catNames);
    
    ylim([-0.1 0.05]);
    ylabel('delta r value')
    box off
    title(Figtitle_Inside);
    
end %iPhase