%look at impact of previous trials on Cue and PE

%Looking at the average firing rate for a given window in each of 4
%current/previous reward conditions

clear all;
load ('R_2R.mat');
load ('RAW.mat');

%run linear model and stats? 1 is yes, 0 is no
runanalysis=1;

%divide neurons up by region
NAneurons=strcmp(R_2R.Ninfo(:,3),'NA');
VPneurons=strcmp(R_2R.Ninfo(:,3),'VP');

%get parameters
trialsback=6; %how many trials back to look
PEBaseline=R_2R.Param.BinBase+0.5; %For normalizing activity
CueBaseline=[-11 -1];
CueWindow=[0 0.5]; %what window of activity is analyzed
PEWindow=[-0.6 0.7];
BinDura=R_2R.Param.BinDura;
bins=R_2R.Param.bins;
binint=R_2R.Param.binint;
binstart=R_2R.Param.binstart;


%sorting bin -- which bin the neurons' activity is sorted on for heatmap(in seconds)
SortBinTime=1; %seconds
SortBin=(((SortBinTime-BinDura(2)/2)-binstart)/binint); %convert to bin name

%reset
NN=0;
PEEvMeanz=0;


if runanalysis==1  
    for i=1:length(RAW) %loops through sessions
        if strcmp('NA',RAW(i).Type(1:2)) | strcmp('VP',RAW(i).Type(1:2)) %only look at suc v mal sessions
            %events
            EV3=strmatch('CueP1',RAW(i).Einfo(:,2),'exact');
            EV4=strmatch('CueP2',RAW(i).Einfo(:,2),'exact');
            EV1=strmatch('PEP1',RAW(i).Einfo(:,2),'exact');
            EV2=strmatch('PEP2',RAW(i).Einfo(:,2),'exact');
            Cue=strmatch('Cue',RAW(i).Einfo(:,2),'exact');
            PE=strmatch('PECue',RAW(i).Einfo(:,2),'exact');
            R1=strmatch('Reward1Deliv',RAW(i).Einfo(:,2),'exact');
            R2=strmatch('Reward2Deliv',RAW(i).Einfo(:,2),'exact');

            %% linear model for impact of previous rewards
            %reset
            Xcue=[];
            Xpe=[];
            Y=[];

            %set up the matrix with outcome identity for each session
            rewards1=cat(2,RAW(i).Erast{R1,1}(:,1),ones(length(RAW(i).Erast{R1,1}(:,1)),1));
            rewards2=cat(2,RAW(i).Erast{R2,1}(:,1),zeros(length(RAW(i).Erast{R2,1}(:,1)),1));
            rewards=cat(1,rewards1,rewards2);
            [rewards(:,1),ind]=sort(rewards(:,1));
            rewards(:,2)=rewards(ind,2);

            %cue
            firstcue=find(RAW(i).Erast{Cue,1}==min(RAW(i).Erast{Cue,1}(RAW(i).Erast{Cue,1}(:,1)>rewards(trialsback),1))); %find first cue with at least 5 rewards prior
            for k=firstcue+1:length(RAW(i).Erast{Cue,1}(:,1))
                time=RAW(i).Erast{Cue,1}(k,1);
                entry=find(rewards==max(rewards(rewards(:,1)<time)));
                for m=1:trialsback
                    Xcue(k-trialsback,m)=rewards(entry+1-m,2);
                end
            end

            %PE
            for k=trialsback+1:length(RAW(i).Erast{PE,1}(:,1))
                time=RAW(i).Erast{PE,1}(k,1);
                entry=find(rewards==max(rewards(rewards(:,1)<time)));
                for m=1:trialsback
                    Xpe(k-trialsback,m)=rewards(entry+1-m,2);
                end
            end    

            for j= 1:length(RAW(i).Nrast) %Number of neurons within sessions

                NN=NN+1;

                %cue
                rewspk=0;
                basespk=0;

                %get mean baseline firing for all cues
                [Bcell1,B1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{Cue},CueBaseline,{2});% makes trial by trial rasters for baseline
                for y= 1:B1n
                    basespk(1,y)=sum(Bcell1{1,y}>CueBaseline(1));
                end

                Bhz=basespk/(CueBaseline(1,2)-CueBaseline(1,1));
                Bmean=nanmean(Bhz);
                Bstd=nanstd(Bhz);

                %get trial by trial firing rate for all cue trials
                [EVcell,EVn]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{Cue},CueWindow,{2});% makes trial by trial rasters for event
                for y= 1:EVn
                    rewspk(y,1)=sum(EVcell{1,y}>CueWindow(1));
                end       
                Y=((rewspk(trialsback+1:end,1)/(CueWindow(1,2)-CueWindow(1,1)))-Bmean)/Bstd; %normalize the activity to baseline

                %true data
                linmdl{NN,1}=fitlm(Xcue,Y);
                R_2R.RewHist.LinMdlCueWeights(NN,1:trialsback)=linmdl{NN,1}.Coefficients.Estimate(2:trialsback+1)';
                R_2R.RewHist.LinMdlCuePVal(NN,1:trialsback)=linmdl{NN,1}.Coefficients.pValue(2:trialsback+1)';

                %shuffled
                YSh=Y(randperm(length(Y)));
                linmdlSh{NN,1}=fitlm(Xcue,YSh);
                R_2R.RewHist.LinMdlCueWeightsSh(NN,1:trialsback)=linmdlSh{NN,1}.Coefficients.Estimate(2:trialsback+1)';
                R_2R.RewHist.LinMdlCuePValSh(NN,1:trialsback)=linmdlSh{NN,1}.Coefficients.pValue(2:trialsback+1)';

                %PE
                rewspk=0;
                basespk=0;

                %get mean baseline firing for all PEs
                [Bcell1,B1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{PE},PEBaseline,{2});% makes trial by trial rasters for baseline
                for y= 1:B1n
                    basespk(1,y)=sum(Bcell1{1,y}>PEBaseline(1));
                end

                Bhz=basespk/(PEBaseline(1,2)-PEBaseline(1,1));
                Bmean=nanmean(Bhz);
                Bstd=nanstd(Bhz);

                %get trial by trial firing rate for all PE trials
                [EVcell,EVn]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{PE},PEWindow,{2});% makes trial by trial rasters for event
                for y= 1:EVn
                    rewspk(y,1)=sum(EVcell{1,y}>PEWindow(1));
                end       
                Y=((rewspk(trialsback+1:end,1)/(PEWindow(1,2)-PEWindow(1,1)))-Bmean)/Bstd; %normalize the activity to baseline

                %true data
                linmdl{NN,1}=fitlm(Xpe,Y);
                R_2R.RewHist.LinMdlPEWeights(NN,1:trialsback)=linmdl{NN,1}.Coefficients.Estimate(2:trialsback+1)';
                R_2R.RewHist.LinMdlPEPVal(NN,1:trialsback)=linmdl{NN,1}.Coefficients.pValue(2:trialsback+1)';

                %shuffled
                YSh=Y(randperm(length(Y)));
                linmdlSh{NN,1}=fitlm(Xpe,YSh);
                R_2R.RewHist.LinMdlPEWeightsSh(NN,1:trialsback)=linmdlSh{NN,1}.Coefficients.Estimate(2:trialsback+1)';
                R_2R.RewHist.LinMdlPEPValSh(NN,1:trialsback)=linmdlSh{NN,1}.Coefficients.pValue(2:trialsback+1)';  

                %% stats comparing effect of current and previous reward on PE
                %resetting
                Bcell=0;
                EV1spk=0;
                EV2spk=0;
                EV1norm=0;
                EV2norm=0;

                %get mean baseline firing for all PE trials
                [Bcell1,B1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV1},PEBaseline,{2});% makes trial by trial rasters for baseline
                for y= 1:B1n
                    Bcell(1,y)=sum(Bcell1{1,y}>PEBaseline(1));
                end
                [Bcell2,B2n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV2},PEBaseline,{2});% makes trial by trial rasters for baseline
                for z= 1:B2n
                    Bcell(1,z+B1n)=sum(Bcell2{1,z}>PEBaseline(1));
                end

                Bhz=Bcell/(PEBaseline(1,2)-PEBaseline(1,1));
                Bmean=nanmean(Bhz);
                Bstd=nanstd(Bhz);

                %get trial by trial firing rate for post suc trials
                [EV1cell,EV1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV1},PEWindow,{2});% makes trial by trial rasters for event
                for y= 1:EV1n
                    EV1spk(1,y)=sum(EV1cell{1,y}>PEWindow(1));
                end       
                EV1hz=EV1spk/(PEWindow(1,2)-PEWindow(1,1));
                for x= 1:EV1n
                    EV1norm(1,x)=((EV1hz(1,x)-Bmean)/Bstd);
                end

                %get trial by trial firing rate for post mal trials
                [EV2cell,EV2n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV2},PEWindow,{2});% makes trial by trial rasters for event
                for y= 1:EV2n
                    EV2spk(1,y)=sum(EV2cell{1,y}>PEWindow(1));
                end       
                EV2hz=EV2spk/(PEWindow(1,2)-PEWindow(1,1));
                for x= 1:EV2n
                    EV2norm(1,x)=((EV2hz(1,x)-Bmean)/Bstd);
                end


                PEEvMeanz(NN,1)=nanmean(EV1norm);
                PEEvMeanz(NN,2)=nanmean(EV2norm);

                %% stats comparing effect of current and previous reward on cue
                %resetting
                Bcell=0;
                EV1spk=0;
                EV2spk=0;
                EV1norm=0;
                EV2norm=0;

                %get mean baseline firing for all cue trials
                [Bcell1,B1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV3},CueBaseline,{2});% makes trial by trial rasters for baseline
                for y= 1:B1n
                    Bcell(1,y)=sum(Bcell1{1,y}>CueBaseline(1));
                end
                [Bcell2,B2n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV4},CueBaseline,{2});% makes trial by trial rasters for baseline
                for z= 1:B2n
                    Bcell(1,z+B1n)=sum(Bcell2{1,z}>CueBaseline(1));
                end

                Bhz=Bcell/(CueBaseline(1,2)-CueBaseline(1,1));
                Bmean=nanmean(Bhz);
                Bstd=nanstd(Bhz);

                %get trial by trial firing rate for post suc trials
                [EV1cell,EV1n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV3},CueWindow,{2});% makes trial by trial rasters for event
                for y= 1:EV1n
                    EV1spk(1,y)=sum(EV1cell{1,y}>CueWindow(1));
                end       
                EV1hz=EV1spk/(CueWindow(1,2)-CueWindow(1,1));
                for x= 1:EV1n
                    EV1norm(1,x)=((EV1hz(1,x)-Bmean)/Bstd);
                end

                %get trial by trial firing rate for post mal trials
                [EV2cell,EV2n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV4},CueWindow,{2});% makes trial by trial rasters for event
                for y= 1:EV2n
                    EV2spk(1,y)=sum(EV2cell{1,y}>CueWindow(1));
                end       
                EV2hz=EV2spk/(CueWindow(1,2)-CueWindow(1,1));
                for x= 1:EV2n
                    EV2norm(1,x)=((EV2hz(1,x)-Bmean)/Bstd);
                end


                CueEvMeanz(NN,1)=nanmean(EV1norm);
                CueEvMeanz(NN,2)=nanmean(EV2norm);

                fprintf('Neuron # %d\n',NN);
            end
        end
        R_2R.RewHist.PrevRewPEMeanz=PEEvMeanz;
        R_2R.RewHist.PrevRewCueMeanz=CueEvMeanz;

    end
end

%% which neurons to look at for stats and plotting?

% Sel=R_2R.SucN | R_2R.MalN; %only look at reward-selective neurons
Sel=NAneurons | VPneurons; %look at all neurons
%Sel=R_2R.RewHist.LinMdlPVal(:,2)<0.1; %only neurons with significant impact of previous trial
%Sel=R_2R.RewHist.LinMdlPEWeights(:,2)<-1; %only neurons with strong negative impact of previous trial

%% ANOVAs

%setup and run ANOVA for effects of current reward, previous reward, and region on reward firing
PrevRew=cat(2,zeros(length(NAneurons),1),ones(length(NAneurons),1));
Region=cat(2,NAneurons,NAneurons);
Rat=cat(2,R_2R.Ninfo(:,4),R_2R.Ninfo(:,4));

%to look at a specific selection of cells
PEEvMeanz=R_2R.RewHist.PrevRewPEMeanz(Sel,:);
CueEvMeanz=R_2R.RewHist.PrevRewCueMeanz(Sel,:);
PrevRew=PrevRew(Sel,:);
Region=Region(Sel,:);
Rat=Rat(Sel,:);

%cue
%each region individually
[~,R_2R.RewHist.PrevRewStatsCueVPSubj{1,1},R_2R.RewHist.PrevRewStatsCueVPSubj{1,2}]=anovan(reshape(CueEvMeanz(VPneurons,:),[sum(VPneurons)*2 1]),{reshape(PrevRew(VPneurons,:),[sum(VPneurons)*2 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*2 1])},'varnames',{'Previous Reward','Rat'},'random',[2],'display','off','model','full');
[~,R_2R.RewHist.PrevRewStatsCueNASubj{1,1},R_2R.RewHist.PrevRewStatsCueNASubj{1,2}]=anovan(reshape(CueEvMeanz(NAneurons,:),[sum(NAneurons)*2 1]),{reshape(PrevRew(NAneurons,:),[sum(NAneurons)*2 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*2 1])},'varnames',{'Previous Reward','Rat'},'random',[2],'display','off','model','full');

%region comparison
[~,R_2R.RewHist.PrevRewStatsCueSubj{1,1},R_2R.RewHist.PrevRewStatsCueSubj{1,2}]=anovan(CueEvMeanz(:),{PrevRew(:),Region(:),Rat(:)},'varnames',{'Previous Reward','Region','Rat'},'random',[3],'nested',[0 0 0;0 0 0;0 1 0],'display','off','model','full');

%pe
%each region individually
[~,R_2R.RewHist.PrevRewStatsPEVPSubj{1,1},R_2R.RewHist.PrevRewStatsPEVPSubj{1,2}]=anovan(reshape(PEEvMeanz(VPneurons,:),[sum(VPneurons)*2 1]),{reshape(PrevRew(VPneurons,:),[sum(VPneurons)*2 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*2 1])},'varnames',{'Previous Reward','Rat'},'random',[2],'display','off','model','full');
[~,R_2R.RewHist.PrevRewStatsPENASubj{1,1},R_2R.RewHist.PrevRewStatsPENASubj{1,2}]=anovan(reshape(PEEvMeanz(NAneurons,:),[sum(NAneurons)*2 1]),{reshape(PrevRew(NAneurons,:),[sum(NAneurons)*2 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*2 1])},'varnames',{'Previous Reward','Rat'},'random',[2],'display','off','model','full');

%region comparison
[~,R_2R.RewHist.PrevRewStatsPESubj{1,1},R_2R.RewHist.PrevRewStatsPESubj{1,2}]=anovan(PEEvMeanz(:),{PrevRew(:),Region(:),Rat(:)},'varnames',{'Previous Reward','Region','Rat'},'random',[3],'nested',[0 0 0;0 0 0;0 1 0],'display','off','model','full');


%setup and run ANOVA for effects of shuffle, trial, and region on coefficient
Trial=[];
Region=[];
Rat=[];
for i=1:trialsback
    Trial=cat(2,Trial,i*ones(length(NAneurons),1));
    Region=cat(2,Region,NAneurons);
    Rat=cat(2,Rat,R_2R.Ninfo(:,4));
end
Trial=cat(2,Trial,Trial);
Region=cat(2,Region,Region);
Rat=cat(2,Rat,Rat);
Shuffd=cat(2,zeros(length(NAneurons),trialsback),ones(length(NAneurons),trialsback));

%Cue
Coeffs=cat(2,R_2R.RewHist.LinMdlCueWeights(:,1:trialsback),R_2R.RewHist.LinMdlCueWeightsSh(:,1:trialsback));
[~,R_2R.RewHist.LinCoeffStatsCueVPSubj{1,1},R_2R.RewHist.LinCoeffStatsCueVPSubj{1,2}]=anovan(reshape(Coeffs(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),{reshape(Shuffd(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),reshape(Trial(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');
[~,R_2R.RewHist.LinCoeffStatsCueNASubj{1,1},R_2R.RewHist.LinCoeffStatsCueNASubj{1,2}]=anovan(reshape(Coeffs(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),{reshape(Shuffd(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),reshape(Trial(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');



%PE
Coeffs=cat(2,R_2R.RewHist.LinMdlPEWeights(:,1:trialsback),R_2R.RewHist.LinMdlPEWeightsSh(:,1:trialsback));
[~,R_2R.RewHist.LinCoeffStatsPEVPSubj{1,1},R_2R.RewHist.LinCoeffStatsPEVPSubj{1,2}]=anovan(reshape(Coeffs(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),{reshape(Shuffd(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),reshape(Trial(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*2*(trialsback) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');
[~,R_2R.RewHist.LinCoeffStatsPENASubj{1,1},R_2R.RewHist.LinCoeffStatsPENASubj{1,2}]=anovan(reshape(Coeffs(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),{reshape(Shuffd(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),reshape(Trial(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*2*(trialsback) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');


%% plotting
XaxisCue=[-0.5 1];
XaxisPE=[-1 2];
Ishow=find(R_2R.Param.Tm>=XaxisCue(1) & R_2R.Param.Tm<=XaxisCue(2));
Ushow=find(R_2R.Param.Tm>=XaxisPE(1) & R_2R.Param.Tm<=XaxisPE(2));
time1=R_2R.Param.Tm(Ishow);
time2=R_2R.Param.Tm(Ushow);

%color map
[magma,inferno,plasma,viridis]=colormaps;
colormap(plasma);
c=[-100 2000];ClimE=sign(c).*abs(c).^(1/4);%colormap

%colors
sucrose=[.95  0.55  0.15];
maltodextrin=[.9  0.3  .9];
water=[0.00  0.75  0.75];
total=[0.3  0.1  0.8];
inh=[0.1 0.021154 0.6];
exc=[0.9 0.75 0.205816];
NAc=[0.5  0.1  0.8];
VP=[0.3  0.7  0.1];

%extra colors to make a gradient
sucrosem=[.98  0.8  0.35];
sucrosel=[1  1  0.4];
maltodextrinm=[1  0.75  1];
maltodextrinl=[1  0.8  1];

RD1P1=strcmp('CueP1', R_2R.Erefnames);
RD1P2=strcmp('CueP2', R_2R.Erefnames);
RD2P1=strcmp('PEP1', R_2R.Erefnames);
RD2P2=strcmp('PEP2', R_2R.Erefnames);

%% Get mean firing according to previous trial and then plot

%NAc

%plot suc after suc
psth1=nanmean(R_2R.Ev(RD1P1).PSTHz(Sel&NAneurons,Ishow),1); 
sem1=nanste(R_2R.Ev(RD1P1).PSTHz(Sel&NAneurons,Ishow),1); %calculate standard error of the mean
up1=psth1+sem1;
down1=psth1-sem1;

%plot suc after malt
psth2=nanmean(R_2R.Ev(RD1P2).PSTHz(Sel&NAneurons,Ishow),1); 
sem2=nanste(R_2R.Ev(RD1P2).PSTHz(Sel&NAneurons,Ishow),1); %calculate standard error of the mean
up2=psth2+sem2;
down2=psth2-sem2;




%make the plot
subplot(2,4,1);
hold on;
title(['Cue response, NAc'])
plot(time1,psth1,'Color',sucrose,'linewidth',1);
plot(time1,psth2,'Color',maltodextrin,'linewidth',1);

patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up1,down1(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);

plot([-2 5],[0 0],':','color','k','linewidth',0.75);
plot([0 0],[-2 8],':','color','k','linewidth',0.75);
plot([CueWindow(1) CueWindow(1)],[-2 8],'color','b','linewidth',0.85);
plot([CueWindow(2) CueWindow(2)],[-2 8],'color','b','linewidth',0.85);
axis([XaxisCue(1) XaxisCue(2) min(down2)-0.15 max(up1)+0.2]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from cue');
legend('Post-suc','Post-mal','location','northeast');

if cell2mat(R_2R.RewHist.PrevRewStatsCueNASubj{1,1}(2,7))<0.05
    plot(CueWindow(1)+(CueWindow(2)-CueWindow(1))/2,max(up2)+0.1,'*','color','k','markersize',13);
end


subplot(2,4,5);
hold on;
title('PE response, NAc');
%plot malt after suc
psth3=nanmean(R_2R.Ev(RD2P1).PSTHz(Sel&NAneurons,Ushow),1); 
sem3=nanste(R_2R.Ev(RD2P1).PSTHz(Sel&NAneurons,Ushow),1); %calculate standard error of the mean
up3=psth3+sem3;
down3=psth3-sem3;

%plot malt after malt
psth4=nanmean(R_2R.Ev(RD2P2).PSTHz(Sel&NAneurons,Ushow),1); 
sem4=nanste(R_2R.Ev(RD2P2).PSTHz(Sel&NAneurons,Ushow),1); %calculate standard error of the mean
up4=psth4+sem4;
down4=psth4-sem4;
plot(time2,psth3,'Color',sucrose,'linewidth',1);
plot(time2,psth4,'Color',maltodextrin,'linewidth',1);
patch([time2,time2(end:-1:1)],[up3,down3(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);
patch([time2,time2(end:-1:1)],[up4,down4(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);

plot([-2 5],[0 0],':','color','k','linewidth',0.75);
plot([0 0],[-2 8],':','color','k','linewidth',0.75);
plot([PEWindow(1) PEWindow(1)],[-2 8],'color','b','linewidth',0.85);
plot([PEWindow(2) PEWindow(2)],[-2 8],'color','b','linewidth',0.85);
axis([XaxisPE(1) XaxisPE(2) min(down3)-0.2 max(up4)+0.2]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from PE');
legend('Post-suc','Post-mal','location','northeast');

if cell2mat(R_2R.RewHist.PrevRewStatsPENASubj{1,1}(2,7))<0.05
    plot(PEWindow(1)+(PEWindow(2)-PEWindow(1))/2,max(up3)+0.1,'*','color','k','markersize',13);
end


%VP

%plot suc after suc
psth1=nanmean(R_2R.Ev(RD1P1).PSTHz(Sel&VPneurons,Ishow),1); 
sem1=nanste(R_2R.Ev(RD1P1).PSTHz(Sel&VPneurons,Ishow),1); %calculate standard error of the mean
up1=psth1+sem1;
down1=psth1-sem1;

%plot suc after malt
psth2=nanmean(R_2R.Ev(RD1P2).PSTHz(Sel&VPneurons,Ishow),1); 
sem2=nanste(R_2R.Ev(RD1P2).PSTHz(Sel&VPneurons,Ishow),1); %calculate standard error of the mean
up2=psth2+sem2;
down2=psth2-sem2;

%make the plot
subplot(2,4,2);
title(['Cue response, VP'])
hold on;
plot(time1,psth1,'Color',sucrose,'linewidth',1);
plot(time1,psth2,'Color',maltodextrin,'linewidth',1);

patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up1,down1(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);

plot([-2 5],[0 0],':','color','k','linewidth',0.75);
plot([0 0],[-2 8],':','color','k','linewidth',0.75);
plot([CueWindow(1) CueWindow(1)],[-2 8],'color','b','linewidth',0.85);
plot([CueWindow(2) CueWindow(2)],[-2 8],'color','b','linewidth',0.85);
axis([XaxisCue(1) XaxisCue(2) min(down2)-0.3 max(up1)+0.3]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from cue');
legend('Post-suc','Post-mal','location','northeast');

if cell2mat(R_2R.RewHist.PrevRewStatsCueVPSubj{1,1}(2,7))<0.05
    plot(CueWindow(1)+(CueWindow(2)-CueWindow(1))/2,max(up2)+0.1,'*','color','k','markersize',13);
end

subplot(2,4,6);
title('PE response, VP');
hold on;
%plot malt after suc
psth3=nanmean(R_2R.Ev(RD2P1).PSTHz(Sel&VPneurons,Ushow),1); 
sem3=nanste(R_2R.Ev(RD2P1).PSTHz(Sel&VPneurons,Ushow),1); %calculate standard error of the mean
up3=psth3+sem3;
down3=psth3-sem3;

%plot malt after malt
psth4=nanmean(R_2R.Ev(RD2P2).PSTHz(Sel&VPneurons,Ushow),1); 
sem4=nanste(R_2R.Ev(RD2P2).PSTHz(Sel&VPneurons,Ushow),1); %calculate standard error of the mean
up4=psth4+sem4;
down4=psth4-sem4;

plot(time2,psth3,'Color',sucrose,'linewidth',1);
plot(time2,psth4,'Color',maltodextrin,'linewidth',1);
patch([time2,time2(end:-1:1)],[up3,down3(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);
patch([time2,time2(end:-1:1)],[up4,down4(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);

plot([-2 5],[0 0],':','color','k','linewidth',0.75);
plot([0 0],[-2 8],':','color','k','linewidth',0.75);
plot([PEWindow(1) PEWindow(1)],[-2 8],'color','b','linewidth',0.85);
plot([PEWindow(2) PEWindow(2)],[-2 8],'color','b','linewidth',0.85);
axis([XaxisPE(1) XaxisPE(2) min(down4)-0.3 max(up3)+0.3]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from PE');
legend('Post-suc','Post-mal','location','northeast');

if cell2mat(R_2R.RewHist.PrevRewStatsPEVPSubj{1,1}(2,7))<0.05
    plot(PEWindow(1)+(PEWindow(2)-PEWindow(1))/2,max(up3)+0.1,'*','color','k','markersize',13);
end


%% plot linear model coefficients

%Plot all neurons
Sel=NAneurons<2;

%coefficients for each trial
subplot(2,4,3);
hold on;
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlCueWeights(Sel&NAneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlCueWeights(Sel&NAneurons,1:trialsback),1),'color',NAc);
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlCueWeights(Sel&VPneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlCueWeights(Sel&VPneurons,1:trialsback),1),'color',VP);
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlCueWeightsSh(Sel&NAneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlCueWeights(Sel&NAneurons,1:trialsback),1),'color','k');
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlCueWeightsSh(Sel&VPneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlCueWeights(Sel&VPneurons,1:trialsback),1),'color','k');
xlabel('Trials back');
ylabel('Mean coefficient weight');
title('Linear model coefficients');
axis([0 trialsback+1 -0.5 1]);
plot([-1 trialsback+1],[0 0],':','color','k','linewidth',0.75);
legend('NAc','VP','Shuff');

%stats to check if VP and NAc are greater than chance
R_2R.RewHist.LinCoeffCueMultComp=[];
[c,~,~,~]=multcompare(R_2R.RewHist.LinCoeffStatsCueNASubj{1,2},'dimension',[1,2],'display','off');
[d,~,~,~]=multcompare(R_2R.RewHist.LinCoeffStatsCueVPSubj{1,2},'dimension',[1,2],'display','off');
for i=1:trialsback

    %NAc vs shuff
    Cel=c(:,1)==2*(i-1)+1 & c(:,2)==2*(i-1)+2;
    if c(Cel,6)<0.05 R_2R.RewHist.LinCoeffCueMultComp(i,1)=1; else R_2R.RewHist.LinCoeffCueMultComp(i,1)=0; end
    R_2R.RewHist.LinCoeffCueMultComp(i,2)=c(Cel,2);

    %VP vs shuff
    Cel=d(:,1)==2*(i-1)+1 & d(:,2)==2*(i-1)+2;
    if d(Cel,6)<0.05 R_2R.RewHist.LinCoeffCueMultComp(i,3)=1; else R_2R.RewHist.LinCoeffCueMultComp(i,3)=0; end
    R_2R.RewHist.LinCoeffCueMultComp(i,4)=d(Cel,4);
end
plot([1:trialsback]-0.1,(R_2R.RewHist.LinCoeffCueMultComp(:,1)-0.5)*1.3,'*','color',NAc); %VP vs shuff
plot([1:trialsback]+0.1,(R_2R.RewHist.LinCoeffCueMultComp(:,3)-0.5)*1.3,'*','color',VP); %NAc vs shuff


%number of neurons with significant weights
subplot(2,4,4);
hold on;
plot(1:trialsback,sum(R_2R.RewHist.LinMdlCuePVal(Sel&NAneurons,1:trialsback)<0.05,1)/sum(Sel&NAneurons),'color',NAc);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlCuePVal(Sel&VPneurons,1:trialsback)<0.05,1)/sum(Sel&VPneurons),'color',VP);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlCuePValSh(Sel&NAneurons,1:trialsback)<0.05,1)/sum(Sel&NAneurons),'color',NAc/3);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlCuePValSh(Sel&VPneurons,1:trialsback)<0.05,1)/sum(Sel&VPneurons),'color',VP/3);
axis([0 trialsback+1 0 0.5]);
xlabel('Trials back');
ylabel('Fraction of the population');
title('Outcome-modulated cue resp');
legend('NAc','VP','Shuff');

%Chi squared stat for each trial
for i=1:trialsback
    [~,R_2R.RewHist.LinMdlCue2all(i,1),R_2R.RewHist.LinMdlCue2all(i,2)]=crosstab(cat(1,R_2R.RewHist.LinMdlCuePVal(Sel,i)<0.05,R_2R.RewHist.LinMdlCuePValSh(Sel,i)<0.05),cat(1,VPneurons,VPneurons+2));
    [~,R_2R.RewHist.LinMdlCue2region(i,1),R_2R.RewHist.LinMdlCue2region(i,2)]=crosstab(R_2R.RewHist.LinMdlCuePVal(Sel,i)<0.05,VPneurons);
end
%plot([1:trialsback]-0.1,(R_2R.RewHist.LinMdlCue2all(:,2)<0.05)-0.52,'*','color',NAc);
plot([1:trialsback],(R_2R.RewHist.LinMdlCue2region(:,2)<0.05&R_2R.RewHist.LinMdlCue2all(:,2)<0.05)-0.52,'*','color','k');

%% PE

%coefficients for each trial
subplot(2,4,7);
hold on;
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlPEWeights(Sel&NAneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlPEWeights(Sel&NAneurons,1:trialsback),1),'color',NAc);
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlPEWeights(Sel&VPneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlPEWeights(Sel&VPneurons,1:trialsback),1),'color',VP);
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlPEWeightsSh(Sel&NAneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlPEWeights(Sel&NAneurons,1:trialsback),1),'color','k');
errorbar(1:trialsback,nanmean(R_2R.RewHist.LinMdlPEWeightsSh(Sel&VPneurons,1:trialsback),1),nanste(R_2R.RewHist.LinMdlPEWeights(Sel&VPneurons,1:trialsback),1),'color','k');
xlabel('Trials back');
ylabel('Mean coefficient weight');
title('Linear model coefficients');
axis([0 trialsback+1 -0.5 1]);
plot([-1 trialsback+1],[0 0],':','color','k','linewidth',0.75);
legend('NAc','VP','Shuff');


%stats to check if VP and NAc are greater than chance
R_2R.RewHist.LinCoeffPEMultComp=[];
[c,~,~,~]=multcompare(R_2R.RewHist.LinCoeffStatsPENASubj{1,2},'dimension',[1,2],'display','off');
[d,~,~,~]=multcompare(R_2R.RewHist.LinCoeffStatsPEVPSubj{1,2},'dimension',[1,2],'display','off');
for i=1:trialsback

    %NAc vs shuff
    Cel=c(:,1)==2*(i-1)+1 & c(:,2)==2*(i-1)+2;
    if c(Cel,6)<0.05 R_2R.RewHist.LinCoeffPEMultComp(i,1)=1; else R_2R.RewHist.LinCoeffPEMultComp(i,1)=0; end
    R_2R.RewHist.LinCoeffPEMultComp(i,2)=c(Cel,2);

    %VP vs shuff
    Cel=d(:,1)==2*(i-1)+1 & d(:,2)==2*(i-1)+2;
    if d(Cel,6)<0.05 R_2R.RewHist.LinCoeffPEMultComp(i,3)=1; else R_2R.RewHist.LinCoeffPEMultComp(i,3)=0; end
    R_2R.RewHist.LinCoeffPEMultComp(i,4)=d(Cel,4);
end
plot([1:trialsback]-0.1,(R_2R.RewHist.LinCoeffPEMultComp(:,1)-0.5)*1.3,'*','color',NAc); %VP vs shuff
plot([1:trialsback]+0.1,(R_2R.RewHist.LinCoeffPEMultComp(:,3)-0.5)*1.3,'*','color',VP); %NAc vs shuff


%number of neurons with significant weights
subplot(2,4,8);
hold on;
plot(1:trialsback,sum(R_2R.RewHist.LinMdlPEPVal(Sel&NAneurons,1:trialsback)<0.05,1)/sum(Sel&NAneurons),'color',NAc);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlPEPVal(Sel&VPneurons,1:trialsback)<0.05,1)/sum(Sel&VPneurons),'color',VP);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlPEPValSh(Sel&NAneurons,1:trialsback)<0.05,1)/sum(Sel&NAneurons),'color',NAc/3);
plot(1:trialsback,sum(R_2R.RewHist.LinMdlPEPValSh(Sel&VPneurons,1:trialsback)<0.05,1)/sum(Sel&VPneurons),'color',VP/3);
axis([0 trialsback+1 0 0.5]);
xlabel('Trials back');
ylabel('Fraction of the population');
title('Outcome-modulated PE resp');
legend('NAc','VP','Shuff');


%Chi squared stat for each trial
for i=1:trialsback
    [~,R_2R.RewHist.LinMdlPEX2all(i,1),R_2R.RewHist.LinMdlPEX2all(i,2)]=crosstab(cat(1,R_2R.RewHist.LinMdlPEPVal(Sel,i)<0.05,R_2R.RewHist.LinMdlPEPValSh(Sel,i)<0.05),cat(1,VPneurons,VPneurons+2));
    [~,R_2R.RewHist.LinMdlPEX2region(i,1),R_2R.RewHist.LinMdlPEX2region(i,2)]=crosstab(R_2R.RewHist.LinMdlPEPVal(Sel,i)<0.05,VPneurons);
end
%plot([0:trialsback]-0.1,(R_2R.RewHist.LinMdlPE2all(:,2)<0.05)-0.52,'*','color',NAc);
plot([1:trialsback],(R_2R.RewHist.LinMdlPEX2region(:,2)<0.05&R_2R.RewHist.LinMdlPEX2all(:,2)<0.05)-0.52,'*','color','k');

%% stats comparing PE coefficient weight in first 2 trials in selective and non-selective neurons in VP
Sel=R_2R.SucN | R_2R.MalN;
NSel=(R_2R.SucN | R_2R.MalN) == 0;
[~,R_2R.RewHist.SelectiveHistory{1,1},R_2R.RewHist.SelectiveHistory{1,2}]=anovan(cat(1,R_2R.RewHist.LinMdlPEWeights(Sel&VPneurons,1),R_2R.RewHist.LinMdlPEWeights(NSel&VPneurons,1),R_2R.RewHist.LinMdlPEWeights(Sel&VPneurons,2),R_2R.RewHist.LinMdlPEWeights(NSel&VPneurons,2)),...
{cat(1,ones(sum(Sel&VPneurons),1),zeros(sum(NSel&VPneurons),1),ones(sum(Sel&VPneurons),1),zeros(sum(NSel&VPneurons),1)),cat(1,ones(sum(VPneurons),1),zeros(sum(VPneurons),1)),...
cat(1,R_2R.Ninfo(Sel&VPneurons,4),R_2R.Ninfo(NSel&VPneurons,4),R_2R.Ninfo(Sel&VPneurons,4),R_2R.Ninfo(NSel&VPneurons,4))},'varnames',{'Selective','Trial','Subject'},'random',3,'model','full','display','off');

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