%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=0;

%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
Baseline=R_2R.Param.BinBase; %For normalizing activity
Window=[0.8 1.3]; %what window of activity is analyzed
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;
EvMeanz=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))
            %events
            EV1=strmatch('RD1P2',RAW(i).Einfo(:,2),'exact');
            EV2=strmatch('RD1P1',RAW(i).Einfo(:,2),'exact');
            EV3=strmatch('RD2P2',RAW(i).Einfo(:,2),'exact');
            EV4=strmatch('RD2P1',RAW(i).Einfo(:,2),'exact');
            RD=strmatch('RD',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
            X=[];
            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);

            for k=trialsback+1:length(RAW(i).Erast{RD,1}(:,1))
                time=RAW(i).Erast{RD,1}(k,1);
                entry=find(rewards(:,1)==time);
                for m=1:trialsback+1
                    X(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;
                rewspk=0;
                basespk=0;

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

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

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

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

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

                %% stats comparing effect of current and previous reward
                %resetting
                Bcell=0;
                EV1spk=0;
                EV2spk=0;
                EV3spk=0;
                EV4spk=0;
                EV1norm=0;
                EV2norm=0;
                EV3norm=0;
                EV4norm=0;

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

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

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

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

                %get trial by trial firing rate for mal/mal trials
                [EV3cell,EV3n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV3},Window,{2});% makes trial by trial rasters for event
                for y= 1:EV3n
                    EV3spk(1,y)=sum(EV3cell{1,y}>Window(1));
                end       
                EV3hz=EV3spk/(Window(1,2)-Window(1,1));
                for x= 1:EV3n
                    EV3norm(1,x)=((EV3hz(1,x)-Bmean)/Bstd);
                end

                %get trial by trial firing rate for mal/suc trials
                [EV4cell,EV4n]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{EV4},Window,{2});% makes trial by trial rasters for event
                for y= 1:EV4n
                    EV4spk(1,y)=sum(EV4cell{1,y}>Window(1));
                end       
                EV4hz=EV4spk/(Window(1,2)-Window(1,1));
                for x= 1:EV4n
                    EV4norm(1,x)=((EV4hz(1,x)-Bmean)/Bstd);
                end

                EvMeanz(NN,1)=nanmean(EV1norm);
                EvMeanz(NN,2)=nanmean(EV2norm);
                EvMeanz(NN,3)=nanmean(EV3norm);
                EvMeanz(NN,4)=nanmean(EV4norm);

                R_2R.RewHist.PrevRewMeanz=EvMeanz;

                fprintf('Neuron # %d\n',NN);
            end
        end

    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.LinMdlWeights(:,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
CurrRew=cat(2,zeros(length(NAneurons),2),ones(length(NAneurons),2));
PrevRew=cat(2,zeros(length(NAneurons),1),ones(length(NAneurons),1),zeros(length(NAneurons),1),ones(length(NAneurons),1));
Region=cat(2,NAneurons,NAneurons,NAneurons,NAneurons);
Rat=cat(2,R_2R.Ninfo(:,4),R_2R.Ninfo(:,4),R_2R.Ninfo(:,4),R_2R.Ninfo(:,4));

%to look at a specific selection of cells
EvMeanz=R_2R.RewHist.PrevRewMeanz(Sel,:);
CurrRew=CurrRew(Sel,:);
PrevRew=PrevRew(Sel,:);
Region=Region(Sel,:);
Rat=Rat(Sel,:);

%each region individually
[~,R_2R.RewHist.PrevRewStatsVPSubj{1,1},R_2R.RewHist.PrevRewStatsVPSubj{1,2}]=anovan(reshape(EvMeanz(VPneurons,:),[sum(VPneurons)*4 1]),{reshape(CurrRew(VPneurons,:),[sum(VPneurons)*4 1]),reshape(PrevRew(VPneurons,:),[sum(VPneurons)*4 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*4 1])},'varnames',{'Current Reward','Previous Reward','Rat'},'random',[3],'display','off','model','full');
[~,R_2R.RewHist.PrevRewStatsNASubj{1,1},R_2R.RewHist.PrevRewStatsNASubj{1,2}]=anovan(reshape(EvMeanz(NAneurons,:),[sum(NAneurons)*4 1]),{reshape(CurrRew(NAneurons,:),[sum(NAneurons)*4 1]),reshape(PrevRew(NAneurons,:),[sum(NAneurons)*4 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*4 1])},'varnames',{'Current Reward','Previous Reward','Rat'},'random',[3],'display','off','model','full');

%region comparison
[~,R_2R.RewHist.PrevRewStatsSubj{1,1},R_2R.RewHist.PrevRewStatsSubj{1,2}]=anovan(EvMeanz(:),{CurrRew(:),PrevRew(:),Region(:),Rat(:)},'varnames',{'Current Reward','Previous Reward','Region','Rat'},'random',[4],'nested',[0 0 0 0;0 0 0 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+1
    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+1),ones(length(NAneurons),trialsback+1));
Coeffs=cat(2,R_2R.RewHist.LinMdlWeights(:,1:trialsback+1),R_2R.RewHist.LinMdlWeightsSh(:,1:trialsback+1));

[~,R_2R.RewHist.LinCoeffStatsVPSubj{1,1},R_2R.RewHist.LinCoeffStatsVPSubj{1,2}]=anovan(reshape(Coeffs(VPneurons,:),[sum(VPneurons)*2*(trialsback+1) 1]),{reshape(Shuffd(VPneurons,:),[sum(VPneurons)*2*(trialsback+1) 1]),reshape(Trial(VPneurons,:),[sum(VPneurons)*2*(trialsback+1) 1]),reshape(Rat(VPneurons,:),[sum(VPneurons)*2*(trialsback+1) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');
[~,R_2R.RewHist.LinCoeffStatsNASubj{1,1},R_2R.RewHist.LinCoeffStatsNASubj{1,2}]=anovan(reshape(Coeffs(NAneurons,:),[sum(NAneurons)*2*(trialsback+1) 1]),{reshape(Shuffd(NAneurons,:),[sum(NAneurons)*2*(trialsback+1) 1]),reshape(Trial(NAneurons,:),[sum(NAneurons)*2*(trialsback+1) 1]),reshape(Rat(NAneurons,:),[sum(NAneurons)*2*(trialsback+1) 1])},'varnames',{'Shuffd','Trial','Subject'},'random',[3],'display','off','model','full');
%% plotting
Xaxis=[-0.5 3];
Ishow=find(R_2R.Param.Tm>=Xaxis(1) & R_2R.Param.Tm<=Xaxis(2));
time1=R_2R.Param.Tm(Ishow);

%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('RD1P1', R_2R.Erefnames);
RD1P2=strcmp('RD1P2', R_2R.Erefnames);
RD2P1=strcmp('RD2P1', R_2R.Erefnames);
RD2P2=strcmp('RD2P2', 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;


%plot malt after suc
psth3=nanmean(R_2R.Ev(RD2P1).PSTHz(Sel&NAneurons,Ishow),1); 
sem3=nanste(R_2R.Ev(RD2P1).PSTHz(Sel&NAneurons,Ishow),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,Ishow),1); 
sem4=nanste(R_2R.Ev(RD2P2).PSTHz(Sel&NAneurons,Ishow),1); %calculate standard error of the mean
up4=psth4+sem4;
down4=psth4-sem4;

%make the plot
subplot(2,4,1);
hold on;
title(['Reward response, current/prev reward'])
plot(time1,psth2,'Color',sucrosem,'linewidth',1);
plot(time1,psth1,'Color',sucrose,'linewidth',1);
plot(time1,psth4,'Color',maltodextrinm,'linewidth',1);
plot(time1,psth3,'Color',maltodextrin,'linewidth',1);
patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],sucrosem,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up1,down1(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up3,down3(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up4,down4(end:-1:1)],maltodextrinm,'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([Window(1) Window(1)],[-2 8],'color','b','linewidth',0.85);
plot([Window(2) Window(2)],[-2 8],'color','b','linewidth',0.85);
axis([Xaxis(1) Xaxis(2) min(down3)-0.15 max(up2)+0.2]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from RD');
legend('Suc/mal','Suc/suc','Mal/mal','Mal/suc','location','northeast');

if cell2mat(R_2R.RewHist.PrevRewStatsNASubj{1,1}(3,7))<0.05
    plot(Window(1)+(Window(2)-Window(1))/2,max(up2)+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;

%plot malt after suc
psth3=nanmean(R_2R.Ev(RD2P1).PSTHz(Sel&VPneurons,Ishow),1); 
sem3=nanste(R_2R.Ev(RD2P1).PSTHz(Sel&VPneurons,Ishow),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,Ishow),1); 
sem4=nanste(R_2R.Ev(RD2P2).PSTHz(Sel&VPneurons,Ishow),1); %calculate standard error of the mean
up4=psth4+sem4;
down4=psth4-sem4;

%make the plot
subplot(2,4,5);
hold on;
title(['Reward response, current/prev reward'])
plot(time1,psth2,'Color',sucrosem,'linewidth',1);
plot(time1,psth1,'Color',sucrose,'linewidth',1);
plot(time1,psth4,'Color',maltodextrinm,'linewidth',1);
plot(time1,psth3,'Color',maltodextrin,'linewidth',1);
patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],sucrosem,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up1,down1(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up3,down3(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up4,down4(end:-1:1)],maltodextrinm,'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([Window(1) Window(1)],[-2 8],'color','b','linewidth',0.85);
plot([Window(2) Window(2)],[-2 8],'color','b','linewidth',0.85);
axis([Xaxis(1) Xaxis(2) min(down3)-0.3 max(up2)+0.3]);
ylabel('Mean firing (z-score)');
xlabel('Seconds from RD');
legend('Suc/mal','Suc/suc','Mal/mal','Mal/suc','location','northeast');

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

%% Plotting heatmaps of each condition

%NAc

%sucrose responses following sucrose
subplot(2,24,[10 11]); %heatmap of water preferring neurons' response to water
Neurons=R_2R.Ev(RD1P1).PSTHz(Sel&NAneurons,Ishow); %get the firing rates of neurons of interest
SucResp=cat(1,R_2R.BinStatRwrd{SortBin+1,1}.R1Mean); %sucrose responses
TMP=SucResp(Sel&NAneurons); %find the magnitude of the excitations for sucrose for this bin
TMP(isnan(TMP))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimg]=sort(TMP);
Neurons=Neurons(SORTimg,:); %sort the neurons by magnitude
imagesc(time1,[1,sum(Sel&NAneurons,1)],Neurons,ClimE);
title(['Suc after suc']);
xlabel('Seconds from RD');
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%following maltodextrin
subplot(2,24,[7 8]); %heatmap of sucrose preferring neurons' response to maltodextrin
Neurons=R_2R.Ev(RD1P2).PSTHz(Sel&NAneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons same as before
imagesc(time1,[1,sum(Sel&NAneurons,1)],Neurons,ClimE);
title(['Suc after mal']);
ylabel('All neurons plotted individually');
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%maltodextrin responses following sucrose
subplot(2,24,[16 17]); %heatmap of water preferring neurons' response to water
Neurons=R_2R.Ev(RD2P1).PSTHz(Sel&NAneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons as before
imagesc(time1,[1,sum(Sel&NAneurons,1)],Neurons,ClimE);
title(['Mal after suc']);
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%following maltodextrin
subplot(2,24,[13 14]); %heatmap of sucrose preferring neurons' response to maltodextrin
Neurons=R_2R.Ev(RD2P2).PSTHz(Sel&NAneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons same as before
imagesc(time1,[1,sum(Sel&NAneurons,1)],Neurons,ClimE);
title(['Mal after mal']);
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%VP

%sucrose responses following sucrose
subplot(2,24,[34 35]); %heatmap of water preferring neurons' response to water
Neurons=R_2R.Ev(RD1P1).PSTHz(Sel&VPneurons,Ishow); %get the firing rates of neurons of interest
SucResp=cat(1,R_2R.BinStatRwrd{SortBin+1,1}.R1Mean); %sucrose responses
TMP=SucResp(Sel&VPneurons); %find the magnitude of the excitations for sucrose for this bin
TMP(isnan(TMP))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimg]=sort(TMP);
Neurons=Neurons(SORTimg,:); %sort the neurons by magnitude
imagesc(time1,[1,sum(Sel&VPneurons,1)],Neurons,ClimE);
title(['Suc after suc']);
xlabel('Seconds from RD');
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%following maltodextrin
subplot(2,24,[31 32]); %heatmap of sucrose preferring neurons' response to maltodextrin
Neurons=R_2R.Ev(RD1P2).PSTHz(Sel&VPneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons same as before
imagesc(time1,[1,sum(Sel&VPneurons,1)],Neurons,ClimE);
title(['Suc after mal']);
ylabel('All neurons plotted individually');
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%maltodextrin responses following sucrose
subplot(2,24,[40 41]); %heatmap of water preferring neurons' response to water
Neurons=R_2R.Ev(RD2P1).PSTHz(Sel&VPneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons as before
imagesc(time1,[1,sum(Sel&VPneurons,1)],Neurons,ClimE);
title(['Mal after suc']);
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%following maltodextrin
subplot(2,24,[37 38]); %heatmap of sucrose preferring neurons' response to maltodextrin
Neurons=R_2R.Ev(RD2P2).PSTHz(Sel&VPneurons,Ishow); %get the firing rates of neurons of interest
Neurons=Neurons(SORTimg,:); %sort the neurons same as before
imagesc(time1,[1,sum(Sel&VPneurons,1)],Neurons,ClimE);
title(['Mal after mal']);
hold on;
plot([0 0],[0 sum(Sel)],':','color','k','linewidth',0.75);

%% plot linear model coefficients

%Plot all neurons
Sel=NAneurons<2;

%coefficients for each trial
subplot(2,4,4);
hold on;
errorbar(0:trialsback,nanmean(R_2R.RewHist.LinMdlWeights(Sel&NAneurons,1:trialsback+1),1),nanste(R_2R.RewHist.LinMdlWeights(Sel&NAneurons,1:trialsback+1),1),'color',NAc,'linewidth',1);
errorbar(0:trialsback,nanmean(R_2R.RewHist.LinMdlWeights(Sel&VPneurons,1:trialsback+1),1),nanste(R_2R.RewHist.LinMdlWeights(Sel&VPneurons,1:trialsback+1),1),'color',VP,'linewidth',1);
errorbar(0:trialsback,nanmean(R_2R.RewHist.LinMdlWeightsSh(Sel&NAneurons,1:trialsback+1),1),nanste(R_2R.RewHist.LinMdlWeights(Sel&NAneurons,1:trialsback+1),1),'color','k','linewidth',1);
errorbar(0:trialsback,nanmean(R_2R.RewHist.LinMdlWeightsSh(Sel&VPneurons,1:trialsback+1),1),nanste(R_2R.RewHist.LinMdlWeights(Sel&VPneurons,1:trialsback+1),1),'color','k','linewidth',1);
xlabel('Trials back');
ylabel('Mean coefficient weight');
title('Linear model coefficients');
legend('NAc','VP','Shuff');

axis([-1 trialsback+1 -1 2.8]);
plot([-1 trialsback+1],[0 0],':','color','k','linewidth',0.75);

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

    %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.LinCoeffMultComp(i,1)=1; else R_2R.RewHist.LinCoeffMultComp(i,1)=0; end
    R_2R.RewHist.LinCoeffMultComp(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.LinCoeffMultComp(i,3)=1; else R_2R.RewHist.LinCoeffMultComp(i,3)=0; end
    R_2R.RewHist.LinCoeffMultComp(i,4)=d(Cel,4);
end
subplot(2,4,4);
hold on;
plot([0:trialsback]-0.1,(R_2R.RewHist.LinCoeffMultComp(:,1)-0.5)*5,'*','color',NAc); %VP vs shuff
plot([0:trialsback]+0.1,(R_2R.RewHist.LinCoeffMultComp(:,3)-0.5)*5,'*','color',VP); %NAc vs shuff


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

%Chi squared stat for each trial
R_2R.RewHist.LinMdlX2all=[];
R_2R.RewHist.LinMdlX2region=[];

for i=1:trialsback+1
    [~,R_2R.RewHist.LinMdlX2all(i,1),R_2R.RewHist.LinMdlX2all(i,2)]=crosstab(cat(1,R_2R.RewHist.LinMdlPVal(Sel,i)<0.05,R_2R.RewHist.LinMdlPValSh(Sel,i)<0.05),cat(1,VPneurons,VPneurons+2));
    [~,R_2R.RewHist.LinMdlX2region(i,1),R_2R.RewHist.LinMdlX2region(i,2)]=crosstab(R_2R.RewHist.LinMdlPVal(Sel,i)<0.05,VPneurons);
end
%plot([0:trialsback]-0.1,(R_2R.LinMdlX2all(:,2)<0.05)-0.52,'*','color',NAc);
plot([0:trialsback],(R_2R.RewHist.LinMdlX2region(:,2)<0.05&R_2R.RewHist.LinMdlX2all(:,2)<0.05)-0.52,'*','color','k');

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