%% generate Figure 10 figure supplement 2

clear all;clc;
tic
global Dura Baseline Tm Tbase BSIZE Tbin
SAVE_FLAG=1;
BSIZE=0.01;
Dura=[-2 4]; Tm=Dura(1):BSIZE:Dura(2);
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) 0];
postwin=[0 .5];
postwin2=[Dura(1):0.05:Dura(2)]; %bounds should match Dura
Slopebounds=[-.5:.05:.5];
R2Bounds=[0:0.05:0.5];
PvalBounds=[0:0.01:0.5];
XI=[1 -1];
RunRegress=1;
SAVEFIG=1;
%R=[];R.Ninfo={}; 
NN=0;
pre=[]; post=[];
if ~exist('RAW'), load ('RAWextendedtraining.mat'); RAW=RAW; end
if ~exist('R'), load('Rextendedtraining.mat'); R=R; end
if RunRegress==0, load('C'); end 
load('Celltype_extendedTraining.mat')

%%

   for i=1:length(RAW)
        EvInd=strmatch('LeverInsertion',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
        Rind=strmatch('First_LP',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);


if RunRegress
  % PREALLOCATING


    Cshuffle.Lat=NaN(MaxNeur, MaxTrials); % (neurons, trials)
    Cshuffle.pre=NaN(MaxNeur,MaxTrials);  % (neurons, trials)
    Cshuffle.post=NaN(MaxNeur,MaxTrials); % (neurons, trials)
    Cshuffle.postwin=NaN(MaxNeur,MaxTrials,MaxWin); % (neurons, trials, windows)
    
    for i=1:length(RAW) %loops through sessions
        EvInd=strmatch('LeverInsertion',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
        Rind=strmatch('First_LP',RAW(i).Einfo(:,2),'exact');  %finds the index of the event
        
        if  sum(EvInd)~=0 && sum(Rind)~=0
            Dcell=MakePSR04(RAW(i).Erast(Rind),RAW(i).Erast{EvInd},[-1 60],{2,'first'});
            D=cell2mat(Dcell); %inds=find(~isnan(D));
            D(isnan(D))=60; %trials with no response are set to 60 to allow to use them in the correlation
            for j= 1:size(RAW(i).Nrast,1) %Number of neurons within sessions
                NN=NN+1;
                Cshuffle.Lat(NN,1:length(RAW(i).Erast{EvInd}))=D;
                [PSR2,N2]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EvInd},Dura,{2});% makes trial by trail rasters. PSR1 is a cell(neurons, trials)
                for m=1:size(RAW(i).Erast{EvInd},1) %loops through trials
                    Cshuffle.pre(NN,m)=sum(PSR2{1,m}<prewin(2) & PSR2{1,m}>prewin(1))/(abs(prewin(2)-prewin(1))); %neurons, trials
                    Cshuffle.post(NN,m)=sum(PSR2{1,m}<postwin(2) & PSR2{1,m}>postwin(1))/(abs(postwin(2)-postwin(1)));
                    for k=1:(length(postwin2)-1) %loops through time windows.
                        Cshuffle.postwin(NN,m,k)=sum(PSR2{1,m}<postwin2(k+1) & PSR2{1,m}>postwin2(k))/(abs(postwin2(k+1)-postwin2(k)));
                    end %k loop
                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.LatShuffled(NN,:)=shuffle(Cshuffle.Lat(NN,:));
        RegX=[Cshuffle.LatShuffled(NN,:)', ones(size(Cshuffle.LatShuffled,2),1)];
         [Cshuffle.SLOPEpostshuffled(NN,n),Cshuffle.STATSpostshuffled(NN,n)]= corr(Cshuffle.post(NN,:)',Cshuffle.LatShuffled(NN,:)','rows','pairwise','type','Spearman');
    end
    fprintf('Shuffled Iteration # %d\n',n);
    % Saving data in C structure and excel file
    save('Cshuffle.mat', 'Cshuffle');
end
end

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

selection1=R.Structure~=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.Lat(NN,:)', ones(size(Cshuffle.Lat,2),1)];
        [Cshuffle.SLOPEpre,Cshuffle.STATSpre]=corr(Cshuffle.pre(NN,:)',Cshuffle.Lat(NN,:)','rows','pairwise','type','Spearman');
        [Cshuffle.SLOPEpost(NN,:),Cshuffle.STATSpost(NN,:)]= corr(Cshuffle.post(NN,:)',Cshuffle.Lat(NN,:)','rows','pairwise','type','Spearman');
        [Cshuffle.SLOPEpostpre(NN,:),Cshuffle.STATSpostpre(NN,:)]= corr((Cshuffle.post(NN,:)-Cshuffle.pre(NN,:))',Cshuffle.Lat(NN,:)','rows','pairwise','type','Spearman'); 
        fprintf('CORREL Neuron ID # %d\n',NN);
    end
    
    Cshuffle.SelectionLAT(:,1)=R.Structure==10 & R.TRN(:,1)~=0 & Cshuffle.STATSpost(:,1)<PStat;
    Cshuffle.SelectionLAT(:,2)=R.Structure==20 & R.TRN(:,1)~=0 & Cshuffle.STATSpost(:,1)<PStat;
    % Saving data in C structure and excel file
    save('Cshuffle.mat', 'Cshuffle');
 
%% Figure of distributions of means correlation coefficients and number of significant neurons
selection1=R.Structure==10 & Celltype(:,1)==1;
selection2=R.Structure==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/sum(Celltype(:,1)==1)),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');
propSigSh=Cshuffle.SIGshuffled*100/sum(Celltype(:,1)==1);
OTpval(1,1)=(sum(propSigSh>SigNUM1)+1)/(length(propSigSh)+1);
OTpval(2,1)=(sum(propSigSh>SigNUM2)+1)/(length(propSigSh)+1);
% Plot of correlation coefficients from example shuffled data set
subplot(2,2,1);
sel1=R.Structure~=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');
hold on;
xlabel('% correlation coefficient shuffle data') % x-axis label
ylabel('# of neurons')% y-axis label

% %Plot of correlation coefficients with significance marked based on
% %bootstrapping
subplot(2,2,2);
sel1=R.Structure~=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