%% generate figure 4
% Analysis of phasic Start and Stop neurons in the extended training dataset

clc;clear;close all
Task=1;
load('table_OT.mat')
load('Rextendedtraining.mat');
load('RatID_extendedtraining.mat');
load('Celltype_extendedTraining.mat');
load('TRN_DT5_extendedtrain.mat')
%% --- --- Plotting color-coded and average PSTHs --- ---

Yaxis=[-2 6];
Xaxis1=[-0.25 1];
Xaxis2=[-2 0.25];

Ishow1=find(R.Param.Tm>=Xaxis1(1) & R.Param.Tm<=Xaxis1(2));
Ishow2=find(R.Param.Tm>=Xaxis2(1) & R.Param.Tm<=Xaxis2(2));

Eventnames={'';'LI';'';'LP1';'';'LP2';'';'LP3';'';'LP4';'';'LP5';'';'PE'};
namebars=categorical({'All START','EXC pre-LP1'});
labelxaxis={'time from LI','time from LP1'};

% ----------- COLORS ---------
myblue=[0 113/255 188/255];
mypurple=[200/255 0 255/255];
mydarkblue=[0 0/255 255/255];
TickSize=[0.015 0.02];

time1=-0.25:0.01:1;
time2=-2:0.01:0.25;

%% Pie Plot
%proportion of start neurons responding before the 1st lever press

DLSselectionPRELP_start=R.Coord(:,4)==10 & Celltype(:,1)==1 & TRNmatrice(:,1)==1 & table_OT(:,5)==2 & table_OT(:,6)==1;
DLSselectionSTART=R.Coord(:,4)==10 & Celltype(:,1)==1 & table_OT(:,5)==2 & table_OT(:,6)==1;

DMSselectionPRELP_start=R.Coord(:,4)==20 & Celltype(:,1)==1 & TRNmatrice(:,1)==1 & table_OT(:,5)==2 & table_OT(:,6)==1;
DMSselectionSTART=R.Coord(:,4)==20 & Celltype(:,1)==1 & table_OT(:,5)==2 & table_OT(:,6)==1;

DSselectionPRELP_start=R.Coord(:,4)~=0 & Celltype(:,1)==1 & TRNmatrice(:,1)==1 & table_OT(:,5)==2 & table_OT(:,6)==1;

DLS_X(1)=sum(DLSselectionSTART);
DMS_X(1)=sum(DMSselectionSTART);
DLS_X(2)=sum(DLSselectionPRELP_start);
DMS_X(2)=sum(DMSselectionPRELP_start);
DS_X=sum(DSselectionPRELP_start);

DLS_tot(1)=sum(R.Coord(:,4)==10 & Celltype(:,1)==1 & table_OT(:,6)==1);
DMS_tot(1)=sum(R.Coord(:,4)==20 & Celltype(:,1)==1 & table_OT(:,6)==1);
DS_tot=sum(R.Coord(:,4)~=0 & Celltype(:,1)==1 & table_OT(:,6)==1);

DS_Percent(1,:)=DS_X*100./DS_tot;
DS_Percent(2,:)=100-DS_Percent(1,:);

figure
subplot(2,4,1)
pie(DS_Percent)
title('Phasic Start');
legend({'EXC pre-LP1','All Start'})


%% Average PSTH start neurons

nplot=0;
for i=1:2 %loop thru event
    nplot=nplot+1;
    
    selectionDLS=R.Coord(:,4)==10 & Celltype(:,1)==1 & table_OT(:,6)==1;
    selectionDMS=R.Coord(:,4)==20 & Celltype(:,1)==1 & table_OT(:,6)==1;
    
    if i==1
        DSconcat_OTz = R.Ev(7).PSTHz(:,Ishow1);
        Xaxis=Xaxis1;time=time1;
        tickdist=0.25;
    else
        DSconcat_OTz = R.Ev(1).PSTHz(:,Ishow2);
        Xaxis=Xaxis2;time=time2;
        tickdist=0.5;
    end
    
    subplot(2,4,[3+(nplot-1)])
    
    DLSconcat_OT=DSconcat_OTz(selectionDLS,:);
    DMSconcat_OT=DSconcat_OTz(selectionDMS,:);
    
    DLSpsth=nanmean(DLSconcat_OT,1);
    DLSsem=nanste(DLSconcat_OT,1);
    DLSup=DLSpsth+DLSsem;
    DLSdown=DLSpsth-DLSsem;
    hold on;
    patch([time,time(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
    plot(time,DLSpsth,'Color',myblue,'linewidth',1.5);
    
    DMSpsth=nanmean(DMSconcat_OT,1);
    DMSsem=nanste(DMSconcat_OT,1);
    DMSup=DMSpsth+DMSsem;
    DMSdown=DMSpsth-DMSsem;
    patch([time,time(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
    plot(time,DMSpsth,'r','linewidth',1.5);
    
    plot([Xaxis(1) Xaxis(end)], [0 0],'LineStyle',':','Color','k');
    plot([0 0], [Yaxis(1) Yaxis(end)],'LineStyle',':','Color','k');
    axis([Xaxis,Yaxis]);
    set(gca,'XTick',[Xaxis(1):tickdist:Xaxis(end)]);
    title('START neurons');
    ylabel('z-score');
    xlabel(labelxaxis(i));
end
 
%% Average PSTH stop neurons
Xaxis2=[-1 1];
Ishow2=find(R.Param.Tm>=Xaxis2(1) & R.Param.Tm<=Xaxis2(2));
time1=-0.25:0.01:1;
time2=-1:0.01:1;
labelxaxis={'time from LLP','time from PE'};

nplot=0;
for i=1:2 %loop thru event
    nplot=nplot+1;
    
    selectionDLS=R.Coord(:,4)==10 & Celltype(:,1)==1 & table_OT(:,6)==2;
    selectionDMS=R.Coord(:,4)==20 & Celltype(:,1)==1 & table_OT(:,6)==2;
    
    if i==1
        DSconcat_OTz = R.Ev(5).PSTHz(:,Ishow1);
        Xaxis=Xaxis1;time=time1;
        tickdist=0.25;
    else
        DSconcat_OTz = R.Ev(6).PSTHz(:,Ishow2);
        Xaxis=Xaxis2;time=time2;
        tickdist=0.5;
    end
    
    subplot(2,4,[5+(nplot-1)])
    
    DLSconcat_OT=DSconcat_OTz(selectionDLS,:);
    DMSconcat_OT=DSconcat_OTz(selectionDMS,:);
    
    DLSpsth=nanmean(DLSconcat_OT,1);
    DLSsem=nanste(DLSconcat_OT,1);
    DLSup=DLSpsth+DLSsem;
    DLSdown=DLSpsth-DLSsem;
    hold on;
    patch([time,time(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
    plot(time,DLSpsth,'Color',myblue,'linewidth',1.5);
    
    DMSpsth=nanmean(DMSconcat_OT,1);
    DMSsem=nanste(DMSconcat_OT,1);
    DMSup=DMSpsth+DMSsem;
    DMSdown=DMSpsth-DMSsem;
    patch([time,time(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
    plot(time,DMSpsth,'r','linewidth',1.5);
    
    %plot([0 0], [-2 2.5],'LineStyle',':','Color','k');
    plot([Xaxis(1) Xaxis(end)], [0 0],'LineStyle',':','Color','k');
    plot([0 0], [Yaxis(1) Yaxis(end)],'LineStyle',':','Color','k');
    axis([Xaxis,Yaxis]);
    set(gca,'XTick',[Xaxis(1):tickdist:Xaxis(end)]);
    title('STOP neurons');
    ylabel('z-score');
    xlabel(labelxaxis(i));
end

 %% Stop neurons, activity around PE at the end of trials VS during inter-trial intervals

titlelabel={'STOP neurons DLS' 'STOP neurons DMS'};
selectionDLS=R.Coord(:,4)==10 & Celltype(:,1)==1 & table_OT(:,6)==2;
selectionDMS=R.Coord(:,4)==20 & Celltype(:,1)==1 & table_OT(:,6)==2;

Xaxis1=[-1 1]; Yaxis=[-2 6];
Ishow=find(R.Param.Tm>=Xaxis1(1) & R.Param.Tm<=Xaxis1(2));

time=[-1:0.01:1];
colors{1,1}=[0.00  0.75  0.75];
colors{1,2}=[0.3  0.1  0.8];

for i=1:2
    if i==1
        selection=selectionDLS;
    else
        selection=selectionDMS;
    end
    meanPostRew_PE=nanmean(R.Ev(6).PSTHz(selection,Ishow),1);
    stePostRew_PE=nanste(R.Ev(6).PSTHz(selection,Ishow),1);
    meanITI_PE=nanmean(R.Ev(12).PSTHz(selection,Ishow),1);
    steITI_PE=nanste(R.Ev(12).PSTHz(selection,Ishow),1);
    
    PRPE_up=meanPostRew_PE+stePostRew_PE;
    PRPE_down=meanPostRew_PE-stePostRew_PE;
    ITIPE_up=meanITI_PE+steITI_PE;
    ITIPE_down=meanITI_PE-steITI_PE;
    
    subplot(2,4,7+i-1)
    patch([time,time(end:-1:1)],[PRPE_up,PRPE_down(end:-1:1)],colors{1,1},'EdgeColor','none');alpha(0.5);
    hold on
    g1=plot(time,meanPostRew_PE,'Color',colors{1,1},'linewidth',1.5);
    patch([time,time(end:-1:1)],[ITIPE_up,ITIPE_down(end:-1:1)],colors{1,2},'EdgeColor','none');alpha(0.5);
    g2=plot(time,meanITI_PE,'Color',colors{1,2},'linewidth',1.5);
    
    plot([-1 1], [0 0],'LineStyle',':','Color','k');
    plot([0 0],Yaxis,'LineStyle',':','Color','k')
    axis([Xaxis1,Yaxis]);
    title(titlelabel(i));
    set(gca,'XTick',[-1:0.5:1]);
    ylabel('z-score');
    xlabel('time from PE');
    hold off
end