%plotting reward response according to trial history
clear all;
load ('R_3R.mat');
load ('RAW.mat');

Xaxis=[-2 5];
Ishow=find(R_3R.Param.Tm>=Xaxis(1) & R_3R.Param.Tm<=Xaxis(2));
time1=R_3R.Param.Tm(Ishow);

runanalysis=1;

%get parameters
trialsback=6; %how many trials back to look
Baseline=R_3R.Param.BinBase; %For normalizing activity
Window=[0.8 1.3]; %what window of activity is analyzed
BinDura=R_3R.Param.BinDura;
bins=R_3R.Param.bins;
binint=R_3R.Param.binint;
binstart=R_3R.Param.binstart;

%colors
sucrose=[0.7  0.3  0];
maltodextrin=[.9  0.3  .9];
water=[0.00  0.75  0.75];
total=[0.3  0.1  0.8];
exc=[0 113/255 188/255];
inh=[240/255 0 50/255];

%extra colors to make a gradient
sucrosem=[0.8  0.7  0.2];
sucrosel=[1  0.8  0.3];
maltodextrinm=[.9  0.6  1];
maltodextrinl=[1  0.8  1];
waterm=[0.1  0.9  0.9];
waterl=[0.3  1  1];
NAc=[0.5  0.1  0.8];
VP=[0.3  0.7  0.1];

RD1P1=strcmp('RD1P1', R_3R.Erefnames);
RD1P2=strcmp('RD1P2', R_3R.Erefnames);
RD1P3=strcmp('RD1P3', R_3R.Erefnames);
RD2P1=strcmp('RD2P1', R_3R.Erefnames);
RD2P2=strcmp('RD2P2', R_3R.Erefnames);
RD2P3=strcmp('RD2P3', R_3R.Erefnames);
RD3P1=strcmp('RD3P1', R_3R.Erefnames);
RD3P2=strcmp('RD3P2', R_3R.Erefnames);
RD3P3=strcmp('RD3P3', R_3R.Erefnames);

Sel = (R_3R.Ev(RD1P1).RespDir<2); %all neurons

%% Plotting sucrose according to previous trial

subplot(2,3,1);
hold on;
title(['Suc response'])
%plot suc preferring neurons to suc
psthE=nanmean(R_3R.Ev(RD1P1).PSTHz(Sel,Ishow),1); 
semE=nanste(R_3R.Ev(RD1P1).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
upE=psthE+semE;
downE=psthE-semE;
plot(time1,psthE,'Color',sucrose,'linewidth',1);

%plot suc preferring neurons to malt
psth2=nanmean(R_3R.Ev(RD1P2).PSTHz(Sel,Ishow),1); 
sem2=nanste(R_3R.Ev(RD1P2).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up2=psth2+sem2;
down2=psth2-sem2;
plot(time1,psth2,'Color',sucrosem,'linewidth',1);

%plot suc preferring neurons to water
psth3=nanmean(R_3R.Ev(RD1P3).PSTHz(Sel,Ishow),1); 
sem3=nanste(R_3R.Ev(RD1P3).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up3=psth3+sem3;
down3=psth3-sem3;
plot(time1,psth3,'Color',sucrosel,'linewidth',1);
patch([time1,time1(end:-1:1)],[upE,downE(end:-1:1)],sucrose,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],sucrosem,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up3,down3(end:-1:1)],sucrosel,'EdgeColor','none');alpha(0.5);
plot([-2 5],[0 0],':','color','k');
plot([0 0],[-2 8],':','color','k');
axis([-2 5 -1 2.5]);
ylabel('Z-score');
legend('Suc after suc','Suc after mal','Suc after wat');
xlabel('Seconds from RD');
%% Plotting maltodextrin according to previous trial

subplot(2,3,2);
hold on;
title(['Mal response'])
%plot mal preferring neurons to suc
psthE=nanmean(R_3R.Ev(RD2P1).PSTHz(Sel,Ishow),1); 
semE=nanste(R_3R.Ev(RD2P1).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
upE=psthE+semE;
downE=psthE-semE;
plot(time1,psthE,'Color',maltodextrin,'linewidth',1);

%plot mal preferring neurons to malt
psth2=nanmean(R_3R.Ev(RD2P2).PSTHz(Sel,Ishow),1); 
sem2=nanste(R_3R.Ev(RD2P2).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up2=psth2+sem2;
down2=psth2-sem2;
plot(time1,psth2,'Color',maltodextrinm,'linewidth',1);

%plot mal preferring neurons to water
psth3=nanmean(R_3R.Ev(RD2P3).PSTHz(Sel,Ishow),1); 
sem3=nanste(R_3R.Ev(RD2P3).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up3=psth3+sem3;
down3=psth3-sem3;
plot(time1,psth3,'Color',maltodextrinl,'linewidth',1);
patch([time1,time1(end:-1:1)],[upE,downE(end:-1:1)],maltodextrin,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],maltodextrinm,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up3,down3(end:-1:1)],maltodextrinl,'EdgeColor','none');alpha(0.5);
plot([-2 5],[0 0],':','color','k');
plot([0 0],[-2 8],':','color','k');
axis([-2 5 -1 2.5]);
ylabel('Z-score');
legend('Mal after suc','Mal after mal','Mal after wat');
xlabel('Seconds from RD');
%% Plotting water according to previous trial

subplot(2,3,3);
hold on;

title(['Water response'])
%plot wat preferring neurons to suc
psthE=nanmean(R_3R.Ev(RD3P1).PSTHz(Sel,Ishow),1); 
semE=nanste(R_3R.Ev(RD3P1).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
upE=psthE+semE;
downE=psthE-semE;
plot(time1,psthE,'Color',water,'linewidth',1);

%plot wat preferring neurons to malt
psth2=nanmean(R_3R.Ev(RD3P2).PSTHz(Sel,Ishow),1); 
sem2=nanste(R_3R.Ev(RD3P2).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up2=psth2+sem2;
down2=psth2-sem2;
plot(time1,psth2,'Color',waterm,'linewidth',1);

%plot wat preferring neurons to water
psth3=nanmean(R_3R.Ev(RD3P3).PSTHz(Sel,Ishow),1); 
sem3=nanste(R_3R.Ev(RD3P3).PSTHz(Sel,Ishow),1); %calculate standard error of the mean
up3=psth3+sem3;
down3=psth3-sem3;
plot(time1,psth3,'Color',waterl,'linewidth',1);
patch([time1,time1(end:-1:1)],[upE,downE(end:-1:1)],water,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up2,down2(end:-1:1)],waterm,'EdgeColor','none');alpha(0.5);
patch([time1,time1(end:-1:1)],[up3,down3(end:-1:1)],waterl,'EdgeColor','none');alpha(0.5);
plot([-2 5],[0 0],':','color','k');
plot([0 0],[-2 8],':','color','k');
axis([-2 5 -1 2.5]);
ylabel('Z-score');
legend('Wat after suc','Wat after mal','Wat after wat');
xlabel('Seconds from RD');

            %% linear model for impact of previous rewards


%reset
NN=0;
EvMeanz=0;

if runanalysis==1  
    for i=1:length(RAW) %loops through sessions
        if strcmp('TH',RAW(i).Type(1:2))
            
            %events
            RD=strmatch('RD',RAW(i).Einfo(:,2),'exact');
            R1=strmatch('Reward1Deliv',RAW(i).Einfo(:,2),'exact');
            R2=strmatch('Reward2Deliv',RAW(i).Einfo(:,2),'exact');
            R3=strmatch('Reward3Deliv',RAW(i).Einfo(:,2),'exact');


            %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));
            rewards3=cat(2,RAW(i).Erast{R3,1}(:,1),-1*ones(length(RAW(i).Erast{R3,1}(:,1)),1));
            rewards=cat(1,rewards1,rewards2,rewards3);
            [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);%,'CategoricalVars',[1:trialsback+1]);
                R_3R.RewHist.LinMdlWeights(NN,1:trialsback+1)=linmdl{NN,1}.Coefficients.Estimate(2:trialsback+2)';
                R_3R.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);%,'CategoricalVars',[1:trialsback+1]);
                R_3R.RewHist.LinMdlWeightsSh(NN,1:trialsback+1)=linmdlSh{NN,1}.Coefficients.Estimate(2:trialsback+2)';
                R_3R.RewHist.LinMdlPValSh(NN,1:trialsback+1)=linmdlSh{NN,1}.Coefficients.pValue(2:trialsback+2)';  
                
                fprintf('Neuron # %d\n',NN);

             end
        end
    end
end

%% plot linear model coefficients

%Plot all neurons
Sel = (R_3R.Ev(RD1P1).RespDir<2); %all neurons

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

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

% %stats to check if VP is greater than NAc
% [c,~,~,~]=multcompare(R_3R.RewHist.LinCoeffStats{1,2},'dimension',[1,2,3],'display','off');
% for i=1:trialsback+1
% 
%     %NAc vs VP
%     Cel=c(:,1)==4*(i-1)+1 & c(:,2)==4*(i-1)+3;
%     if c(Cel,6)<0.05 R_3R.RewHist.LinCoeffMultComp(i,1)=1; else R_3R.RewHist.LinCoeffMultComp(i,1)=0; end
%     R_3R.RewHist.LinCoeffMultComp(i,2)=c(Cel,6);
% 
%     %VP vs shuff
%     Cel=c(:,1)==4*(i-1)+1 & c(:,2)==4*(i-1)+2;
%     if c(Cel,6)<0.05 R_3R.RewHist.LinCoeffMultComp(i,3)=1; else R_3R.RewHist.LinCoeffMultComp(i,3)=0; end
%     R_3R.RewHist.LinCoeffMultComp(i,4)=c(Cel,6);
% 
%     %NAc vs shuff
%     Cel=c(:,1)==4*(i-1)+3 & c(:,2)==4*(i-1)+4;
%     if c(Cel,6)<0.05 R_3R.RewHist.LinCoeffMultComp(i,5)=1; else R_3R.RewHist.LinCoeffMultComp(i,5)=0; end
%     R_3R.RewHist.LinCoeffMultComp(i,6)=c(Cel,6);
% end
% subplot(2,4,4);
% hold on;
% plot([0:trialsback]-0.15,(R_3R.RewHist.LinCoeffMultComp(:,3)-0.5)*5,'*','color',VP); %VP vs shuff
% plot([0:trialsback],(R_3R.RewHist.LinCoeffMultComp(:,5)-0.5)*5,'*','color',NAc); %NAc vs shuff
% plot([0:trialsback]+0.15,(R_3R.RewHist.LinCoeffMultComp(:,1)-0.5)*5,'*','color','k'); %VP vs NAc


%number of neurons with significant weights
subplot(2,3,5);
hold on;
plot(0:trialsback,sum(R_3R.RewHist.LinMdlPVal(Sel,1:trialsback+1)<0.05,1)/sum(Sel),'color',VP,'linewidth',1);
plot(0:trialsback,sum(R_3R.RewHist.LinMdlPValSh(Sel,1:trialsback+1)<0.05,1)/sum(Sel),'color',VP/3,'linewidth',1);
axis([-1 trialsback+1 0 1]);
xlabel('Trials back');
ylabel('Fraction of the population');
title('Reward-modulated neurons');
legend('True','Shuff');

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