%% generate Figure 10 and Figure 10 supplement 1 - Correlation analysis

%firing analysis post lever retraction and correl with PE latency
%Analysis on all neurons without DLS/DMS dissociation
%histogram for correlated neurons in each regions, in each class


clear all;clc;close all
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];
prewinPE=[-0.5 0];
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;
XI=1;
RunRegress=1;
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');



if RunRegress
  % PREALLOCATING
    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);
    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)
    
    
    %
    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;
                C.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. 
                [PSR3,N3]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{Rind},Dura,{2});% ref= post-reward PE / 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);
                    C.prePE(NN,m)=sum(PSR2{1,m}<prewin(2) & PSR2{1,m}>prewin(1))/(abs(prewin(2)-prewin(1))); %neurons, trials
                    C.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
    for NN=1:MaxNeur
        RegX=[C.LatPE(NN,:)', ones(size(C.LatPE,2),1)];
        [C.SLOPEprePE,C.STATSprePE]=corr(C.prePE(NN,:)',C.LatPE(NN,:)','rows','pairwise','type','Spearman');
        [C.SLOPEpostPE1(NN,:),C.STATSpostPE1(NN,:)]= corr(C.postPE1(NN,:)',C.LatPE(NN,:)','rows','pairwise','type','Spearman');
        [C.SLOPEpostprePE(NN,:),C.STATSpostprePE(NN,:)]= corr((C.postPE(NN,:)-C.prePE(NN,:))',C.LatPE(NN,:)','rows','pairwise','type','Spearman');
        fprintf('CORREL Neuron ID # %d\n',NN);
    end
    
    save('C.mat', 'C');
end

%% --------------Plots distribution of correlation variables SINGLE WINDOW -------------
% POST WINDOW
figure(1);

for i=1:2
    sel1=Coord(:,4)==i*10;
    sel2= sel1 & C.STATSpostPE1(:,1)<PStat;  
    h=subplot(2,2,i);
    xmin=floor(min(C.SLOPEpostPE1(sel1,1)));
    xmax=ceil(max(C.SLOPEpostPE1(sel1,1)));
    step=0.05;
    xcenters=xmin:step:xmax;
    Nsel1=hist(C.SLOPEpostPE1(sel1,1),xcenters);
    Nsel2=hist(C.SLOPEpostPE1(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');
    %axis([-10 10 0 20]) %note: their is an outlier @-14 for shell
    pval1=signrank(C.SLOPEpostPE1(sel1,1));
    pval2=signrank(C.SLOPEpostPE1(sel2,1));
    MEANsel1=mean(C.SLOPEpostPE1(sel2,1));
    plot([MEANsel1 MEANsel1], [0 10],'Color','r');
    %text(-5.8,9.8,sprintf('p=%d', pval1), 'Parent',h);
    %text(-5.8,8.8,sprintf('p=%d', pval2), 'Parent',h);
    subplot(2,2,i+2)
    [N,xcenters2]=hist(C.STATSpostPE1(sel2,1),20);
    bar(xcenters2, N, 'k', 'EdgeColor', 'k', 'BarWidth', 1); alpha(0.3);
    axis ([0 ceil(max(xcenters2)*10)/10 0 ceil(max(N))+1]);
end

save('C.mat', 'C');