%% generate Figure 10 figure supplement 3

clear all;clc;
tic
global Dura Baseline Tm Tbase BSIZE Tbin
SAVE_FLAG=1;
BSIZE=0.05;
Dura=[-25 60]; Tm=-1:BSIZE:60;
Baseline=[-2 0]; Tbase=Baseline(1):BSIZE:Baseline(2); %used to calculate Z-scores
Tbin=-1:0.005:1; %window used to determine the optimal binsize
PStat=0.05;
prewin=[Dura(1) -20];
postwin=[0 .5];
EventWin=[-0.5 0.5];
postwin2=[Dura(1):0.05:Dura(2)]; %bounds should match Dura
Slopebounds=[-1.5:0.25:1.5];
R2Bounds=[0:0.05:0.5];
PvalBounds=[0:0.01:0.5];
Struct=[10 20];
XI=1;
RunRegress=0;
SAVEFIG=1;
%R=[];R.Ninfo={}; 
NN=0;
pre=[]; post=[];
if ~exist('RAW'), load ('RAWextendedtraining.mat'); RAW=RAW; end
if ~exist('Ev'), load('Rextendedtraining_light.mat'); end
load('C.mat');
load('Celltype_extendedTraining.mat')
% 1 for inclusion of omission in analysis, max time=60s / 0 for exclusion omission

%Smoothing
Smoothing=1; %0 for raw and 1 for smoothing
SmoothTYPE='lowess';
SmoothSPAN=5;
if Smoothing~=1, SmoothTYPE='NoSmoothing';SmoothSPAN=NaN; end


%%

 % PREALLOCATING
    for i=1:length(RAW)
        EvInd=strmatch('First_LP',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
        Rind=strmatch('ReliableLP_DT',RAW(i).Einfo(:,2),'exact');  %finds the index of the event
        
        if  sum(EvInd)~=0 && sum(Rind)~=0
            DS(i)=length(RAW(i).Erast{EvInd});
            Neur(i)=length(RAW(i).Nrast);
        end
    end
    MaxTrials=max(DS); MaxNeur=sum(Neur); MaxWin=length(postwin);
    Cshuffle.preRRRR=NaN(MaxNeur,MaxTrials);  % (neurons, trials)
    Cshuffle.SEQ=NaN(MaxNeur,MaxTrials); % (neurons, trials)
    Cshuffle.postwinRR=NaN(MaxNeur,MaxTrials,MaxWin); % (neurons, trials, windows)
    
    for i=1:length(RAW) %loops through sessions
        EvInd=strmatch('First_LP',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
        Rind=strmatch('ReliableLP_DT',RAW(i).Einfo(:,2),'exact');  %finds the index of the event
        
        if  sum(EvInd)~=0 && sum(Rind)~=0            
            Dcell2=MakePSR05(RAW(i).Erast(Rind),RAW(i).Erast{EvInd},[-1 60],{2}); % makes trial by trail 
                    
            for j= 1:size(RAW(i).Nrast,1) %Number of neurons within sessions
                NN=NN+1;

                [PSRNeur,Nneur]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EvInd},[-1 60],{1});% collapsed raster
                [PSR2,N2]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EvInd},Dura,{2});% makes trial by trail rasters. PSR1 is a cell(neurons, trials)
                
                [PTHneur,BW1,~]=MakePTH07(PSRNeur,repmat(N2, size(RAW(i).Nrast{j},1),1),{2,0,BSIZE});
                PTHneur=smooth(PTHneur,SmoothSPAN,SmoothTYPE)';               
                Cshuffle.PSTHraw(NN,1:length(Tm))=PTHneur;
                Search=find(Tm>=EventWin(1) & Tm<=EventWin(2));
                [Cshuffle.MaxVal(NN,1),MaxInd]=max(Cshuffle.PSTHraw(NN,Search));
                Cshuffle.MaxTime(NN,1)=Tm(Search(1)+MaxInd-1);
                
                for m=1:size(RAW(i).Erast{EvInd},1) %loops through trials
                    SequenceDur=RAW(i).Eint{1,6}(m)-RAW(i).Eint{1,5}(m);
                    Cshuffle.RR(NN,m)=length(Dcell2{1,m})/SequenceDur;
                    Cshuffle.preRR(NN,m)=sum(PSR2{1,m}<prewin(2) & PSR2{1,m}>prewin(1))/(abs(prewin(2)-prewin(1))); %neurons, trials
                    Cshuffle.SEQ(NN,m)=sum(PSR2{1,m}<SequenceDur & PSR2{1,m}>0)/SequenceDur;
                end %m loop
                fprintf('Neuron ID # %d\n',NN);
            end %j loop
        end %if
    end %i loop
    
    
    
    
    %% Runs the regression analysis on XX number of iterations of shuffled data
    
XX=1000;

for n=1:XX
    for NN=1:MaxNeur
        Cshuffle.SEQShuffled(NN,:)=shuffle(Cshuffle.SEQ(NN,:));
        RegX=[Cshuffle.SEQShuffled(NN,:)', ones(size(Cshuffle.SEQShuffled,2),1)];
        [Cshuffle.SLOPEpostshuffled(NN,n),Cshuffle.STATSpostshuffled(NN,n)]= corr(Cshuffle.RR(NN,:)',Cshuffle.SEQShuffled(NN,:)','rows','pairwise','type','Spearman');
    end
    fprintf('Shuffled Iteration # %d\n',n);
    save('Cshuffle.mat', 'Cshuffle');
end


%% Calculates number of significant correlations and the mean correlation coefficient for each 

selection1=Coord(:,4)~=0;

for q=1:length(Cshuffle.SLOPEpostshuffled(1,:))
    Cshuffle.SLOPEmeanshuffled(q)=nanmean(Cshuffle.SLOPEpostshuffled(selection1,q));
    Cshuffle.SIGshuffled(q)=sum(Cshuffle.STATSpostshuffled(selection1,q)<0.05);
    Cshuffle.NEGSIGshuffled(q)=sum(Cshuffle.STATSpostshuffled(selection1,q)<0.05 & Cshuffle.SLOPEpostshuffled(selection1,q)<0);
end



%% Calculates correlation coefficients and significant values for real data


    for NN=1:MaxNeur
        RegX=[Cshuffle.SEQ(NN,:)', ones(size(Cshuffle.SEQ,2),1)];
        [Cshuffle.SLOPEpre,Cshuffle.STATSpre]=corr(Cshuffle.preRR(NN,:)',Cshuffle.SEQ(NN,:)','rows','pairwise','type','Spearman');
        [Cshuffle.SLOPEpost(NN,:),Cshuffle.STATSpost(NN,:)]= corr(Cshuffle.RR(NN,:)',Cshuffle.SEQ(NN,:)','rows','pairwise','type','Spearman');
        [Cshuffle.SLOPEpostpre(NN,:),Cshuffle.STATSpostpre(NN,:)]= corr((Cshuffle.RR(NN,:)-Cshuffle.preRR(NN,:))',Cshuffle.SEQ(NN,:)','rows','pairwise','type','Spearman');
        fprintf('CORREL Neuron ID # %d\n',NN);
    end 
    save('Cshuffle.mat', 'Cshuffle');
 
%% Figure of distributions of means correlation coefficients and number of significant neurons
selection1=Coord(:,4)==10 & Celltype(:,1)==1;
selection2=Coord(:,4)==20 & Celltype(:,1)==1;
figure;
% Plot of distribution of mean correlation coefficients for all of the
% shuffled data
subplot(2,2,3);
histogram(Cshuffle.SLOPEmeanshuffled,100);
hold on;
MEANsel1=nanmean(Cshuffle.SLOPEpost(selection1,1));
MEANsel2=nanmean(Cshuffle.SLOPEpost(selection2,1));
plot([MEANsel1 MEANsel1], [0 50],'Color','b');
plot([MEANsel2 MEANsel2], [0 50],'Color','r');
xlabel('Mean correlation coefficient') % x-axis label
ylabel('# of iterations')% y-axis label
%Plot of distribution of number of significantly correlated neurons across
%all the shuffled data sets
subplot(2,2,4);
histogram((Cshuffle.SIGshuffled*100/849),40);
xlabel('% of significantly correlated unit') % x-axis label
ylabel('# of iterations')% y-axis label
hold on;

SigNUM1=100*sum(Cshuffle.STATSpost(selection1,1)<0.05)/(sum(selection1));
SigNUM2=100*sum(Cshuffle.STATSpost(selection2,1)<0.05)/(sum(selection2));
plot([SigNUM1 SigNUM1], [0 100],'Color','b');
plot([SigNUM2 SigNUM2], [0 100],'Color','r');
% Plot of correlation coefficients from example shuffled data set
subplot(2,2,1);
sel1=Coord(:,4)~=0;
sel2=sel1 & Cshuffle.STATSpostshuffled(:,1)<PStat;
xmin=floor(min(Cshuffle.SLOPEpostshuffled(sel1,1)));
xmax=ceil(max(Cshuffle.SLOPEpostshuffled(sel1,1)));
step=0.05;
xcenters=xmin:step:xmax;
Nsel1=hist(Cshuffle.SLOPEpostshuffled(sel1,1),xcenters);
Nsel2=hist(Cshuffle.SLOPEpostshuffled(sel2,1),xcenters);
hold on;
bar(xcenters, Nsel1, 'w', 'EdgeColor', 'k', 'BarWidth', 1);alpha(0.3);
bar(xcenters, Nsel2, 'k', 'EdgeColor', 'k', 'BarWidth', 1);
plot([0 0], [0 30],'LineStyle',':','Color','k');
MEANsel1=mean(Cshuffle.SLOPEpostshuffled(sel1,1));
plot([MEANsel1 MEANsel1], [0 10],'Color','r');
xlabel('% correlation coefficient shuffle data') % x-axis label
ylabel('# of neurons')% y-axis label
hold on;

% %Plot of correlation coefficients with significance marked based on
% %bootstrapping
subplot(2,2,2);
sel1=Coord(:,4)~=0;
sel2= sel1 & Cshuffle.STATSpost(:,1)<PStat;
xmin=floor(min(Cshuffle.SLOPEpost(sel1,1)));
xmax=ceil(max(Cshuffle.SLOPEpost(sel1,1)));
step=0.05;
xcenters=xmin:step:xmax;
Nsel1=hist(Cshuffle.SLOPEpost(sel1,1),xcenters);
Nsel2=hist(Cshuffle.SLOPEpost(sel2,1),xcenters);
hold on;
bar(xcenters, Nsel1, 'w', 'EdgeColor', 'k', 'BarWidth', 1);alpha(0.3);
bar(xcenters, Nsel2, 'k', 'EdgeColor', 'k', 'BarWidth', 1);
plot([0 0], [0 30],'LineStyle',':','Color','k');
MEANsel1=mean(Cshuffle.SLOPEpost(sel2,1));
plot([MEANsel1 MEANsel1], [0 10],'Color','r');
xlabel('% correlation coefficient real data') % x-axis label
ylabel('# of neurons')% y-axis label