%% generate Figure 10 figure supplement 4

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')

%path='C:\FRED\GALLO\Exps\2009-present GoNoGo\FINAL\PLX Files GoNoGo\GoNoGo NEX files\RESULTgn14.xls';
%%
for i=1:length(RAW)
    EvInd=strmatch('Last_LP',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
    Rind=strmatch('FirstPE_PostRew',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

  C.LatPE=NaN(MaxNeur, MaxTrials); % (neurons, trials)
  C.prePE=NaN(MaxNeur,MaxTrials);  % (neurons, trials)
  C.postPE=NaN(MaxNeur,MaxTrials); % (neurons, trials)
  C.postwinPE=NaN(MaxNeur,MaxTrials,MaxWin); % (neurons, trials, windows)
  
  Cshuffle.LatPE=NaN(MaxNeur, MaxTrials); % (neurons, trials)
  Cshuffle.prePE=NaN(MaxNeur,MaxTrials);  % (neurons, trials)
  Cshuffle.postPE=NaN(MaxNeur,MaxTrials); % (neurons, trials)
  Cshuffle.postwinPE=NaN(MaxNeur,MaxTrials,MaxWin); % (neurons, trials, windows)
  
  for i=1:length(RAW) %loops through sessions
      EvInd=strmatch('Last_LP',RAW(i).Einfo(:,2),'exact'); %finds the index of the event
      Rind=strmatch('FirstPE_PostRew',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 10],{2,'first'});
          D=cell2mat(Dcell); %inds=find(~isnan(D));
          D(isnan(D))=60; %trials with no response are set to 10.25 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.LatPE(NN,1:length(RAW(i).Erast{EvInd}))=D;
              [PSR2,N2]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EvInd},Dura,{2});% ref= last lever press / makes trial by trail rasters.
              for m=1:size(RAW(i).Erast{Rind},1) %loops through trials
                  PElatency=RAW(i).Erast{Rind}(m)-RAW(i).Erast{EvInd}(m);
                  Cshuffle.prePE(NN,m)=sum(PSR2{1,m}<prewin(2) & PSR2{1,m}>prewin(1))/(abs(prewin(2)-prewin(1))); %neurons, trials
                  Cshuffle.postPE1(NN,m)=sum(PSR2{1,m}<postwin(2) & PSR2{1,m}>postwin(1))/(abs(postwin(2)-postwin(1))); %When last LP is ref, postwin is 0-0.5s post LP
              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.LatPE(NN,:));
        RegX=[Cshuffle.LatShuffled(NN,:)', ones(size(Cshuffle.LatShuffled,2),1)];
        [Cshuffle.SLOPEpostshuffled(NN,n),Cshuffle.STATSpostshuffled(NN,n)]= corr(Cshuffle.postPE1(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.LatPE(NN,:)', ones(size(Cshuffle.LatPE,2),1)];
        [Cshuffle.SLOPEpre,Cshuffle.STATSpre]=corr(Cshuffle.prePE(NN,:)',Cshuffle.LatPE(NN,:)','rows','pairwise','type','Spearman');
        [Cshuffle.SLOPEpost(NN,:),Cshuffle.STATSpost(NN,:)]= corr(Cshuffle.postPE1(NN,:)',Cshuffle.LatPE(NN,:)','rows','pairwise','type','Spearman');  
        [Cshuffle.SLOPEpostpre(NN,:),Cshuffle.STATSpostpre(NN,:)]= corr((Cshuffle.postPE1(NN,:)-Cshuffle.prePE(NN,:))',Cshuffle.LatPE(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');
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=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