%% generate Figure 2.
% Plotting Heatmaps all neurons 1st/last session early training and  extended training
% need R_light datafiles, Celltype datafiles to only include MSN, and TRN
% datafiles to restrict analysis on TRN.


clc;clear;close all
Task=1;
load('Rextendedtraining_light.mat');
load('Celltype_extendedTraining.mat');
load('TRN_DT5_extendedtrain.mat')

%% --- --- Plotting color-coded and average PSTHs --- ---
BrainRegion=[10 20]; %10 for DLS, 20 for DMS
ExcInh=1;        % 1 for excitations, 2 for inhibitions -- Used to find Exc or Inh onsets in R.Ev(.).Onsets

Pstat=0.05;
stepsize=0.3;      % Used for the scatter plot projections

% For COLOR-CODED PSTH only (not Average PSTH)
Norm=0;          % 0 = no norm - max accross all the selected neurons
                 % 1 each neuron is normalized wh its own max
                 % 2 for baseline substraction

Sorting=4;       % 1 to sort neurons by onset
                 % 2 to sort neurons by duration
                 % 3 to sort neurons by PeakTime
                 % 4 to sort by class of neuron
%for RExt ONLY!  % 5 to sort by the magnitude of the cue response
%for Cues ONLY!  % 6 to sort by the magnitude of the RExt response

Xaxis=[0 357];Yaxis=[-1 3];
Xaxis1=[-0.25 0.25];
Ishow=find(Tm>=Xaxis1(1) & Tm<=Xaxis1(2));

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

Eventnames={'';'LI';'';'LP1';'';'LP2';'';'LP3';'';'LP4';'';'LP5';'';'PE'};

time=-0.25:0.25:3.25;time2=1:1:357;

%% Indiv Neuron
%analysis restricted on putative MSN. nonTRN included

%concat with normalized data
DSconcat_OT = Ev(7).PSTHz(:,Ishow);
for i=1:6
    DSconcat_OT = cat(2,DSconcat_OT,Ev(i).PSTHz(:,Ishow));
end

DLSselection=Coord(:,4)==10 & Celltype(:,1)==1 & TRN(:,1)~=0;
DMSselection=Coord(:,4)==20 & Celltype(:,1)==1 & TRN(:,1)~=0;

LPwin=[52:280];
binLP=mean(DSconcat_OT(:,LPwin),2);
SEQwin=[26:331];
binSEQ=mean(DSconcat_OT(TRN(:,1)~=0,SEQwin),2);
TRN(TRN(:,1)~=0,2)=sign(binSEQ);

DLSconcat_OT=DSconcat_OT(DLSselection,:);
DMSconcat_OT=DSconcat_OT(DMSselection,:);

%TMPdls=R.Class(DLSselection,1); TMPdms=R.Class(DMSselection,1);
TMPdls=binLP(DLSselection,1); 
TMPdms=binLP(DMSselection,1);
TMPdls(isnan(TMPdls))=0;
TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimgDLS]=sort(TMPdls);[~,SORTimgDMS]=sort(TMPdms);

DLSconcat_OT=DLSconcat_OT(SORTimgDLS,:);
DMSconcat_OT=DMSconcat_OT(SORTimgDMS,:);


subplot(9,8,[29 30 37 38 45 46 53 54 61 62 69 70])
Clim=[-5 5];
imagesc(time,[1 size(DLSconcat_OT,1)],DLSconcat_OT,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([0.25 0.25],[size(DLSconcat_OT,1)+0.5 0.5],'k')
plot([0.75 0.75],[size(DLSconcat_OT,1)+0.5 0.5],'k')
plot([1.25 1.25],[size(DLSconcat_OT,1)+0.5 0.5],'k')
plot([1.75 1.75],[size(DLSconcat_OT,1)+0.5 0.5],'k')
plot([2.25 2.25],[size(DLSconcat_OT,1)+0.5 0.5],'k')
plot([2.75 2.75],[size(DLSconcat_OT,1)+0.5 0.5],'k')
set(gca,'XTick',[0:0.25:3.25]);
set(gca,'xticklabel',Eventnames)
title('Extended training DLS');
ylabel('neurons');
hold off

subplot(9,8,[31 32 39 40 47 48 55 56 63 64 71 72])
imagesc(time,[1 size(DMSconcat_OT,1)],DMSconcat_OT,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([0.25 0.25],[size(DMSconcat_OT,1)+0.5 0.5],'k')
plot([0.75 0.75],[size(DMSconcat_OT,1)+0.5 0.5],'k')
plot([1.25 1.25],[size(DMSconcat_OT,1)+0.5 0.5],'k')
plot([1.75 1.75],[size(DMSconcat_OT,1)+0.5 0.5],'k')
plot([2.25 2.25],[size(DMSconcat_OT,1)+0.5 0.5],'k')
plot([2.75 2.75],[size(DMSconcat_OT,1)+0.5 0.5],'k')
set(gca,'XTick',[0:0.25:3.25]);
set(gca,'xticklabel',Eventnames)
title('Extended training DMS');
ylabel('neurons');
hold off


%% Average OT PSTH

subplot(9,8,[23 24])

DLSpsth=nanmean(DLSconcat_OT,1);
DLSsem=nanste(DLSconcat_OT,1);
DLSup=DLSpsth+DLSsem;
DLSdown=DLSpsth-DLSsem;
hold on;
patch([time2,time2(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
plot(time2,DLSpsth,'Color',myblue,'linewidth',1.5);

DMSpsth=nanmean(DMSconcat_OT,1);
DMSsem=nanste(DMSconcat_OT,1);
DMSup=DMSpsth+DMSsem;
DMSdown=DMSpsth-DMSsem;
patch([time2,time2(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
plot(time2,DMSpsth,'r','linewidth',1.5);

%plot([0 0], [-2 2.5],'LineStyle',':','Color','k');
plot([0 357], [0 0],'LineStyle',':','Color','k');
plot([51 51],Yaxis,'LineStyle',':','Color','k')
plot([102 102],Yaxis,'LineStyle',':','Color','k')
plot([153 153],Yaxis,'LineStyle',':','Color','k')
plot([204 204],Yaxis,'LineStyle',':','Color','k')
plot([255 255],Yaxis,'LineStyle',':','Color','k')
plot([306 306],Yaxis,'LineStyle',':','Color','k')
axis([Xaxis,Yaxis]);
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
title('Extended training');
ylabel('z-score');

%% Pie Plot
%proportion nonTRN excited inhibited
%analysis restricted on putative MSN

DLSselectionNonTRN=Coord(:,4)==10 & TRN(:,1)==0 & Celltype(:,1)==1;
DLSselectionTRN_excited=Coord(:,4)==10 & TRN(:,1)~=0 & TRN(:,2)>0 & Celltype(:,1)==1;
DLSselectionTRN_Inhibited=Coord(:,4)==10 & TRN(:,1)~=0 & TRN(:,2)<0 & Celltype(:,1)==1;

DMSselectionNonTRN=Coord(:,4)==20 & TRN(:,1)==0 & Celltype(:,1)==1;
DMSselectionTRN_excited=Coord(:,4)==20 & TRN(:,1)~=0 & TRN(:,2)>0 & Celltype(:,1)==1;
DMSselectionTRN_Inhibited=Coord(:,4)==20 & TRN(:,1)~=0 & TRN(:,2)<0 & Celltype(:,1)==1;

DLS_X(1)=sum(DLSselectionTRN_excited);
DLS_X(2)=sum(DLSselectionTRN_Inhibited);
DLS_X(3)=sum(DLSselectionNonTRN);

DMS_X(1)=sum(DMSselectionTRN_excited);
DMS_X(2)=sum(DMSselectionTRN_Inhibited);
DMS_X(3)=sum(DMSselectionNonTRN);

DLS_Percent=DLS_X*100./(sum(Coord(:,4)==10 & Celltype(:,1)==1));
DMS_Percent=DMS_X*100./(sum(Coord(:,4)==20 & Celltype(:,1)==1));

labels={'EXC','INH','non-TRN'};
subplot(9,8,[11 12])
bar(sum(DLS_Percent),'FaceColor',[0.9290 0.6940 0.1250]);
hold on
bar(sum(DLS_Percent(1:2)),'FaceColor',[0.8500 0.3250 0.0980]);
bar(DLS_Percent(1),'FaceColor',[0 0.4470 0.7410]);
axis([0.5 10.5 0 100 ])

subplot(9,8,[19 20])
bar(sum(DMS_Percent),'FaceColor',[0.9290 0.6940 0.1250]);
hold on
bar(sum(DMS_Percent(1:2)),'FaceColor',[0.8500 0.3250 0.0980]);
bar(DMS_Percent(1),'FaceColor',[0 0.4470 0.7410]);
axis([0.5 10.5 0 100 ])


 save('Rextendedtraining_light.mat','Coord','Erefnames','Ev','Ninfo','Tm','TRN')

%% ACQUISITION DATASET

clear Ev 
load('Rearlytraining_light.mat');
load('Celltype_earlyTraining.mat');
load('TRN_DT5_earlytrain.mat')

for i=1:max(Celltype(:,1))
    selectionIN=Ses(i).Coord(:,4)~=0;
    selectsession=Celltype(:,1)==i;

    Ses(i).Celltype(1:length(selectionIN),1)=0;
    Ses(i).Celltype(selectionIN,1)=Celltype(selectsession,2);  
end


%analysis restricted on putative MSN. nonTRN included
%concat with normalized data
titlegraph={'Early training 1st session','Early training last session'};

for j=1:2
    vars={'TMPdls','TMPdms','binSEQ','binLP','DLSselection','DMSselection','DSconcat_ACQ','DLSconcat_ACQ','DMSconcat_ACQ'};
    clear(vars{:});
    
    if j==1
        k=4;
    else
        k=13;
    end
    
    DSconcat_ACQ = Ses(k).Ev(7).PSTHz(:,Ishow);
    for i=1:6
        DSconcat_ACQ = cat(2,DSconcat_ACQ,Ses(k).Ev(i).PSTHz(:,Ishow));
    end
    
    DLSselection=Ses(k).Coord(:,4)==10 & Ses(k).Celltype(:,1)==1 & sesTRN(k).TRN(:,1)~=0;
    DMSselection=Ses(k).Coord(:,4)==20 & Ses(k).Celltype(:,1)==1 & sesTRN(k).TRN(:,1)~=0;
    
    LPwin=[52:280];
    binLP=mean(DSconcat_ACQ(:,LPwin),2);
    SEQwin=[26:331];
    binSEQ=mean(DSconcat_ACQ(:,SEQwin),2);
    sesTRN(k).TRN(:,2)=sign(binSEQ);
    
    DLSconcat_ACQ=DSconcat_ACQ(DLSselection,:);
    DMSconcat_ACQ=DSconcat_ACQ(DMSselection,:);
    
    %TMPdls=R.Class(DLSselection,1); TMPdms=R.Class(DMSselection,1);
    TMPdls=binLP(DLSselection,1);
    TMPdms=binLP(DMSselection,1);
    TMPdls(isnan(TMPdls))=0;
    TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimgDLS]=sort(TMPdls);[~,SORTimgDMS]=sort(TMPdms);
    
    DLSconcat_ACQ=DLSconcat_ACQ(SORTimgDLS,:);
    DMSconcat_ACQ=DMSconcat_ACQ(SORTimgDMS,:);
    
    
    subplot(9,8,[25+(j-1)*24 26+(j-1)*24 33+(j-1)*24 34+(j-1)*24 41+(j-1)*24 42+(j-1)*24])
    Clim=[-5 5];
    imagesc(time,[1 size(DLSconcat_ACQ,1)],DLSconcat_ACQ,Clim); %colorbar;axis tight;
    colormap(jet);
    hold on
    plot([0.25 0.25],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    plot([0.75 0.75],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    plot([1.25 1.25],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    plot([1.75 1.75],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    plot([2.25 2.25],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    plot([2.75 2.75],[size(DLSconcat_ACQ,1)+0.5 0.5],'k')
    set(gca,'XTick',[0:0.25:3.25]);
    set(gca,'xticklabel',Eventnames)
    title('DLS');
    ylabel('neurons');
    hold off

    subplot(9,8,[27+(j-1)*24 28+(j-1)*24 35+(j-1)*24 36+(j-1)*24 43+(j-1)*24 44+(j-1)*24])
    imagesc(time,[1 size(DMSconcat_ACQ,1)],DMSconcat_ACQ,Clim); %colorbar;axis tight;
    colormap(jet);
    hold on
    plot([0.25 0.25],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    plot([0.75 0.75],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    plot([1.25 1.25],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    plot([1.75 1.75],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    plot([2.25 2.25],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    plot([2.75 2.75],[size(DMSconcat_ACQ,1)+0.5 0.5],'k')
    set(gca,'XTick',[0:0.25:3.25]);
    set(gca,'xticklabel',Eventnames)
    title('DMS');
    ylabel('neurons');
    hold off
    
    
    %% Average OT PSTH
    
    subplot(9,8,[7+(j-1)*8 8+(j-1)*8])
    
    DLSpsth=nanmean(DLSconcat_ACQ,1);
    DLSsem=nanste(DLSconcat_ACQ,1);
    DLSup=DLSpsth+DLSsem;
    DLSdown=DLSpsth-DLSsem;
    hold on;
    patch([time2,time2(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
    plot(time2,DLSpsth,'Color',myblue,'linewidth',1.5);
    
    DMSpsth=nanmean(DMSconcat_ACQ,1);
    DMSsem=nanste(DMSconcat_ACQ,1);
    DMSup=DMSpsth+DMSsem;
    DMSdown=DMSpsth-DMSsem;
    patch([time2,time2(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
    plot(time2,DMSpsth,'r','linewidth',1.5);
    
    %plot([0 0], [-2 2.5],'LineStyle',':','Color','k');
    plot([0 357], [0 0],'LineStyle',':','Color','k');
    plot([51 51],Yaxis,'LineStyle',':','Color','k')
    plot([102 102],Yaxis,'LineStyle',':','Color','k')
    plot([153 153],Yaxis,'LineStyle',':','Color','k')
    plot([204 204],Yaxis,'LineStyle',':','Color','k')
    plot([255 255],Yaxis,'LineStyle',':','Color','k')
    plot([306 306],Yaxis,'LineStyle',':','Color','k')
    axis([Xaxis,Yaxis]);
    set(gca,'XTick',[0:25.5:357]);
    set(gca,'xticklabel',Eventnames)
    title(titlegraph(j));
    ylabel('z-score');
    
end

%% Pie Plot
%proportion nonTRN excited inhibited
%analysis restricted on putative MSN

for i=4:13
    
    DSconcat_ACQ = Ses(i).Ev(7).PSTHz(:,Ishow);
    for j=1:6
        DSconcat_ACQ = cat(2,DSconcat_ACQ,Ses(i).Ev(j).PSTHz(:,Ishow));
    end
    SEQwin=[26:331];
    binSEQ=nanmean(DSconcat_ACQ(:,SEQwin),2);
    sesTRN(i).TRN(:,2)=sign(binSEQ);
    
    DLSselectionNonTRN_ACQ=Ses(i).Coord(:,4)==10 & sesTRN(i).TRN(:,1)==0 & Ses(i).Celltype(:,1)==1;
    DLSselectionTRN_excited_ACQ=Ses(i).Coord(:,4)==10 & sesTRN(i).TRN(:,1)~=0 & sesTRN(i).TRN(:,2)>0 & Ses(i).Celltype(:,1)==1;
    DLSselectionTRN_Inhibited_ACQ=Ses(i).Coord(:,4)==10 & sesTRN(i).TRN(:,1)~=0 & sesTRN(i).TRN(:,2)<0 & Ses(i).Celltype(:,1)==1;
    
    DMSselectionNonTRN_ACQ=Ses(i).Coord(:,4)==20 & sesTRN(i).TRN(:,1)==0 & Ses(i).Celltype(:,1)==1;
    DMSselectionTRN_excited_ACQ=Ses(i).Coord(:,4)==20 & sesTRN(i).TRN(:,1)~=0 & sesTRN(i).TRN(:,2)>0 & Ses(i).Celltype(:,1)==1;
    DMSselectionTRN_Inhibited_ACQ=Ses(i).Coord(:,4)==20 & sesTRN(i).TRN(:,1)~=0 & sesTRN(i).TRN(:,2)<0 & Ses(i).Celltype(:,1)==1;
    
    DLS_X_ACQ(i-3,1)=sum(DLSselectionTRN_excited_ACQ);
    DLS_X_ACQ(i-3,2)=sum(DLSselectionTRN_Inhibited_ACQ);
    DLS_X_ACQ(i-3,3)=sum(DLSselectionNonTRN_ACQ);
    
    DMS_X_ACQ(i-3,1)=sum(DMSselectionTRN_excited_ACQ);
    DMS_X_ACQ(i-3,2)=sum(DMSselectionTRN_Inhibited_ACQ);
    DMS_X_ACQ(i-3,3)=sum(DMSselectionNonTRN_ACQ);
    
    DLS_Percent_ACQ(i-3,:)=DLS_X_ACQ(i-3,:)*100./(sum(Ses(i).Coord(:,4)==10 & Ses(i).Celltype(:,1)==1));
    DMS_Percent_ACQ(i-3,:)=DMS_X_ACQ(i-3,:)*100./(sum(Ses(i).Coord(:,4)==20 & Ses(i).Celltype(:,1)==1));
end

labels={'EXC','INH','non-TRN'};
subplot(9,8,[9 10])
bar(DLS_Percent_ACQ,1,'stacked');
axis([0.5 10.5 0 100 ])
subplot(9,8,[17 18])
bar(DMS_Percent_ACQ,1,'stacked');
axis([0.5 10.5 0 100 ])

save('Rearlytraining_light.mat','Erefnames','Ses','Tm') 
save('TRN_DT5_earlytrain.mat','sesTRN')