%% generate figure 3, 5 and Figure 3 - supplement figure 1

% fourrier analysis - hierarchical clustering - Classification approach on
% extended training dataset
% statistics for the extended training dataset

clc;clear;close all;
load('Rextendedtraining_light.mat');
load('RatID_extendedtraining.mat'); %% column 1 = rat ID, column 2 = habit ID; 0=goal-directed / 1=habit
load('Celltype_extendedTraining.mat');
load('TRN_DT5_extendedtrain.mat')

timeconcat=[1:1:357];
BrainRegion=[10 20];
Erefnames=Erefnames(1:7);
Xaxis1=[-0.25 0.25]; Yaxis=[-4 8];
Xaxis2=[1 357]; Yaxis2=[-2 6];
Ishow=find(Tm>=Xaxis1(1) & Tm<=Xaxis1(2));
Eventnames={'';'LI';'';'LP1';'';'LP2';'';'LP3';'';'LP4';'';'LP5';'';'PE'};

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


%% preparation of data
%normalization based on max
for k=1:length(Erefnames)
    Ev(k).PSTH_norm1=normalize(Ev(k).PSTHraw,Ev(k).PSTHrawBL,1);
end

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

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

selection=Coord(:,4)~=0 & Celltype(:,1)==1;% & TRN(:,1)~=0; %exclude NaN values
selectionID=find(Coord(:,4)~=0 & Celltype(:,1)==1);% & TRN(:,1)~=0);
outID=find(Coord(:,4)==0 | Celltype(:,1)~=1);% & TRN(:,1)==0);
DSconcat_OTshort=DSconcat_OTz(selection,:);


%% Clustering on phasic and non phasic neurons 

close all;
f_low=1;f_high=4;

%fourier analysis
[X_all,f_all,P_all,Nt_all]=myfft(DSconcat_OTshort',0.01,1);
total_power=sum(P_all);
P_all_norm = bsxfun(@rdivide,P_all, total_power);


a=find(f_all>f_low);b=find(f_all<f_high);
lower_window_sum_all=sum(P_all_norm(1:a(1),:));
upper_window_sum_all=sum(P_all_norm(a(1):b(end),:));

ratio_all=lower_window_sum_all./upper_window_sum_all;

%hierarchical clustering
X_ratio=[[lower_window_sum_all]' [upper_window_sum_all]']; %#ok<NBRAK>
eva_Ch_ratio = evalclusters(X_ratio,'linkage','CalinskiHarabasz','KList',[1:10]);
c_ratio=eva_Ch_ratio.OptimalY; % if wat to use optimal number of cluster

 Z = linkage(X_ratio,'ward'); % X is the data set, neuron in rows, features in columns
 c = cluster(Z,'Maxclust',3); % determine the number cluster here, if you don't use optimalK
 D = pdist(X_ratio);
leafOrder = optimalleaforder(Z,D);
cutoff = median([Z(end-2,3) Z(end-1,3)]);

 phasicID(selectionID,1)=c_ratio;
 phasicID(outID,1)=NaN;


%% heatmap results clustering
figure(1)

for i=1:2
    ind=find(phasicID==i);
    t=num2str(i);
    Clim=[-5 5];
    DSconcat_OT_class=[];SORTimg=[];MaxVal=[];MeanVal=[]; MaxTime=[];TMP=[];NNid=[];
    DSconcat_OT_class=DSconcat_OT(ind,:);
    
    for NN=1:size(DSconcat_OT_class,1)
        MeanVal(NN,1)=mean(DSconcat_OT_class(NN,52:306),2);
        [MaxVal(NN,1),MaxInd]=max(DSconcat_OT_class(NN,:));
        MaxTime(NN,1)=timeconcat(timeconcat(1)+MaxInd-1);
        if i==2
            MaxTimePhasic=MaxTime;
        end
    end
    if i==1
        TMP=MeanVal;
        subplot(5,3,[4 7 10 13])
    elseif i==2 
        TMP=MeanVal;
        subplot(5,3,[5 8 11 14])
    end
    TMP(isnan(TMP))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimg]=sort(TMP);
    %NNid=cell2mat(Ninfo(phasic_neuron_selection,3));NNid2=NNid(SORTimg);
    DSconcat_OT_class=DSconcat_OT_class(SORTimg,:);
    
    
    imagesc(timeconcat,[1 size(DSconcat_OT_class,1)],DSconcat_OT_class); %colorbar;axis tight;
    colormap(jet);
    hold on
    plot([51 51],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    plot([102 102],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    plot([153 153],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    plot([204 204],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    plot([255 255],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    plot([306 306],[size(DSconcat_OT_class,1)+0.5 0.5],'k')
    set(gca,'XTick',[0:25.5:357]);
    set(gca,'xticklabel',Eventnames)
    hold off
end

%% heatmap phasic versus non phasic neurons

phasicID(phasicID(:,1)==3,1)=1;
subplot(5,3,2)
dendrogram(Z,'ColorThreshold',2, 'Reorder',leafOrder);


for j=1:2 % loop thru structure
    for k=1:max(phasicID(:,1)) % loop thru class
        DSselection(:,1)=Coord(:,4)==BrainRegion(j) & phasicID(:,1)==k;
        DSselectionNumber(j,k)=sum(DSselection);
        DSselectionProp(j,k)=100*sum(DSselection)/sum(Coord(:,4)==BrainRegion(j) & Celltype(:,1)==1);
    end
end
subplot(5,3,3)
X = categorical({'DLS','DMS'});
b2=bar(X,DSselectionProp,'stacked');
ylabel('% neurons');
legend('phasic', 'nonphasic')


%% Figure scatter plot Hierarchical clustering based on fourier

for i=1:length(c_ratio)
    if c_ratio(i,1)==2
        color(i,:)=[255/255 0/255 0/255]; %phasic red
    elseif c_ratio(i,1)==1
        color(i,:)=[0 255/255 0/255]; %sustained green
    end
end
subplot(5,3,1)
scatter(lower_window_sum_all,upper_window_sum_all,[],color)
hold on
title(strcat('hierarchical clustering'))
xlabel('Power low freq') % x-axis label
ylabel('Power int freq')% y-axis label
legend('phasic', 'nonphasic')
hold off

%% Plot heatmap phasic neurons, separation DMS/DLS, 
PhasicNeurons_ID(:,1)=find(phasicID==2);
PhasicNeurons_ID(:,2)=Coord(PhasicNeurons_ID(:,1),4);
PhasicNeurons_ID(:,3)=RatID(PhasicNeurons_ID(:,1),1);

DSconcat_OT_phasic=DSconcat_OTz(PhasicNeurons_ID(:,1),:);
DLSselection=PhasicNeurons_ID(:,2)~=0 & PhasicNeurons_ID(:,2)==10;
DMSselection=PhasicNeurons_ID(:,2)~=0 & PhasicNeurons_ID(:,2)==20;

% sorting DLS
DSconcat_OT_classDLS=DSconcat_OT_phasic(DLSselection,:);
for NN=1:size(DSconcat_OT_classDLS,1)
    [MaxValDLS(NN,1),MaxInd]=max(DSconcat_OT_classDLS(NN,:));
    MaxTimeDLS(NN,1)=timeconcat(timeconcat(1)+MaxInd-1);
end
TMPdls=MaxTimeDLS;
TMPdls(isnan(TMPdls))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimgDLS]=sort(TMPdls);
DSconcat_OT_classDLS=DSconcat_OT_classDLS(SORTimgDLS,:);

% sorting DMS
DSconcat_OT_classDMS=DSconcat_OT_phasic(DMSselection,:);
for NN=1:size(DSconcat_OT_classDMS,1)
    [MaxValDMS(NN,1),MaxInd]=max(DSconcat_OT_classDMS(NN,:));
    MaxTimeDMS(NN,1)=timeconcat(timeconcat(1)+MaxInd-1);
end
TMPdms=MaxTimeDMS;
TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimgDMS]=sort(TMPdms);
DSconcat_OT_classDMS=DSconcat_OT_classDMS(SORTimgDMS,:);

% plot heatmap
figure(2)
subplot(9,4,[1 2 5 6 9 10])
imagesc(timeconcat,[1 size(DSconcat_OT_classDLS,1)],DSconcat_OT_classDLS,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([51 51],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
plot([102 102],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
plot([153 153],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
plot([204 204],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
plot([255 255],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
plot([306 306],[size(DSconcat_OT_classDLS,1)+0.5 0.5],'k')
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
ylabel('DLS neurons')
hold off

subplot(9,4,[17 18 21 22 25 26])
imagesc(timeconcat,[1 size(DSconcat_OT_classDMS,1)],DSconcat_OT_classDMS,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([51 51],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
plot([102 102],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
plot([153 153],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
plot([204 204],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
plot([255 255],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
plot([306 306],[size(DSconcat_OT_classDMS,1)+0.5 0.5],'k')
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
ylabel('DMS neurons')
hold off


%% Hierarchical clustering of phasic neurons based on Sequence-related activity

LIwin=[26:51];%post lever insertion window
LPwin=[52:280];%pre-1stLP to pre-lastLP window
PEwin=[307:332];% pre-PE window

binLI=mean(DSconcat_OT(:,LIwin),2);
binLP=mean(DSconcat_OT(:,LPwin),2);
binPE=mean(DSconcat_OT(:,PEwin),2);
binnedDS=cat(2,binLI,binLP,binPE);   

sel=Coord(:,4)~=0 & Celltype==1 & phasicID==2;
X=binnedDS(sel,:);

Z = linkage(X,'ward');
c = cluster(Z,'Maxclust',4);
eva = evalclusters(X,'linkage','CalinskiHarabasz','KList',[1:10]); %% You can change the criterion here. There are 3 other options. Gap, 'Silhouette' , 'DaviesBouldin'. You can do help evalclusters to use this code.
D = pdist(X);
leafOrder = optimalleaforder(Z,D);
cutoff = median([Z(end-2,3) Z(end-1,3)]);
subplot(9,4,[29 33])
dendrogram(Z,'ColorThreshold',cutoff, 'Reorder',leafOrder);

ClassPHAS(:,1)=eva.OptimalY;
ClassPHAS(:,2:4)=X;


%% eva has the results of fitting the clusters. eva.OptimalK tells us how many clusters fit this data best based on the criterion used.
%% eva.OptimalY tells us the cluster identities of each point.
%% CH criterion clustering labels
ind=[];
DSconcat_OT_Phasic=DSconcat_OTz(sel,:);

for i=1:max(eva.OptimalY)
    DSconcat_OT_classDLS=[]; SORTimgDLS=[];MaxValDLS=[];TMPdls=[];
    DSconcat_OT_classDMS=[]; SORTimgDMS=[];MaxValDMS=[];TMPdms=[];
    ind=find(eva.OptimalY==i);
    
    PhasicNeurons_ID(:,1)=find(phasicID~=1 & ~isnan(phasicID));
    PhasicNeurons_ID(:,2)=Coord(PhasicNeurons_ID(:,1),4);
    PhasicNeurons_ID(:,3)=eva.OptimalY;
    
    DLSselectionP=PhasicNeurons_ID(:,2)==10 &  PhasicNeurons_ID(:,3)==i;
    DMSselectionP=PhasicNeurons_ID(:,2)==20 &  PhasicNeurons_ID(:,3)==i;

    
    % sorting DLS
    DSconcat_OT_classDLS=DSconcat_OT_Phasic(DLSselectionP,:);
    for NN=1:size(DSconcat_OT_classDLS,1)
        MaxValDLS(NN,1)=mean(DSconcat_OT_classDLS(NN,52:306),2);
    end
    TMPdls=MaxValDLS;
    TMPdls(isnan(TMPdls))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimgDLS]=sort(TMPdls);
    DSconcat_OT_classDLS=DSconcat_OT_classDLS(SORTimgDLS,:);
    
    % sorting DMS
    DSconcat_OT_classDMS=DSconcat_OT_Phasic(DMSselectionP,:);
    for NN=1:size(DSconcat_OT_classDMS,1)
        MaxValDMS(NN,1)=mean(DSconcat_OT_classDMS(NN,52:306),2);
    end
    TMPdms=MaxValDMS;
    TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimgDMS]=sort(TMPdms);
    DSconcat_OT_classDMS=DSconcat_OT_classDMS(SORTimgDMS,:);
    
    for j=1:2
        if j==1
            DSconcat_OT_class_DLSDMS=DSconcat_OT_classDLS;
        else
            DSconcat_OT_class_DLSDMS=DSconcat_OT_classDMS;
        end
    subplot(9,4,[3+(i-1)*12+(j-1)*4 4+(i-1)*12+(j-1)*4])
    imagesc(timeconcat,[1 size(DSconcat_OT_class_DLSDMS,1)],DSconcat_OT_class_DLSDMS,Clim);
    colormap(jet);
    hold on
    plot([51 51],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    plot([102 102],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    plot([153 153],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    plot([204 204],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    plot([255 255],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    plot([306 306],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
    set(gca,'XTick',[0:25.5:357]);
    set(gca,'xticklabel',Eventnames)
    hold off
    
    % plot average PSTH
    subplot(9,4,[11+(i-1)*12 12+(i-1)*12])
    DLSpsth=nanmean(DSconcat_OT_classDLS(:,timeconcat),1);
    DLSsem=nanste(DSconcat_OT_classDLS(:,timeconcat),1);
    DLSup=DLSpsth+DLSsem;
    DLSdown=DLSpsth-DLSsem;
    
    DMSpsth=nanmean(DSconcat_OT_classDMS(:,timeconcat),1);
    DMSsem=nanste(DSconcat_OT_classDMS(:,timeconcat),1);
    DMSup=DMSpsth+DMSsem;
    DMSdown=DMSpsth-DMSsem;
    hold on
    patch([timeconcat,timeconcat(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
    g1=plot(timeconcat,DLSpsth,'Color',myblue,'linewidth',1.5);
    patch([timeconcat,timeconcat(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
    g2=plot(timeconcat,DMSpsth,'r','linewidth',1.5);
    
    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([Xaxis2,Yaxis]);
    set(gca,'XTick',[0:25.5:357]);
    set(gca,'xticklabel',Eventnames)
    hold off
    end
end


%% proportion of neurons in each classes / plot stack bars / 2*3 classes with separation phasic / non phasic neurons

DSselection=[];

for j=1:2 % loop thru structure
    for k=1:3 % loop thru class
        DSselection(:,1)=PhasicNeurons_ID(:,2)==BrainRegion(j) & PhasicNeurons_ID(:,3)==k;
        DSselectionNumber(j,k)=sum(DSselection);
        DSselectionProp(j,k)=100*sum(DSselection)/sum(PhasicNeurons_ID(:,2)==BrainRegion(j));
    end
end

subplot(9,4,[30 34])
X=categorical({'DLS','DMS'});
b1=bar(X,DSselectionProp,'stacked');
legend('Start','Stop','Middle')


%% Hierarchical clustering of Nonphasic neurons based on Seqeunce-related activity
figure
sel=Coord(:,4)~=0 & Celltype==1 & phasicID==1;
X=binnedDS(sel,:);

Z = linkage(X,'ward');
c = cluster(Z,'Maxclust',4);
eva = evalclusters(X,'linkage','CalinskiHarabasz','KList',[1:10]); %% You can change the criterion here. There are 3 other options. Gap, 'Silhouette' , 'DaviesBouldin'. You can do help evalclusters to use this code.
D = pdist(X);
leafOrder = optimalleaforder(Z,D);
cutoff = median([Z(end-2,3) Z(end-1,3)]);
subplot(6,4,[17 21])
dendrogram(Z,'ColorThreshold',5, 'Reorder',leafOrder);

ClassNONPHASIC(:,1)=eva.OptimalY;
ClassNONPHASIC(:,2:4)=X;

save('ClassHC.mat','ClassPHAS','ClassNONPHASIC')


%% eva has the results of fitting the clusters. eva.OptimalK tells us how many clusters fit this data best based on the criterion used.
%% eva.OptimalY tells us the cluster identities of each point.
%% CH criterion clustering labels
ind=[];
DSconcat_OT_NonPhasic=DSconcat_OTz(sel,:);
DLSsel=Coord(:,4)==10 & Celltype==1 & phasicID==1;
DMSsel=Coord(:,4)==20 & Celltype==1 & phasicID==1;

% sorting all neurons
DSconcat_OT_NPclassDLS=DSconcat_OTz(DLSsel,:);


for NN=1:size(DSconcat_OT_NPclassDLS,1)
    MeanValdls(NN,1)=mean(DSconcat_OT_NPclassDLS(NN,52:306),2);
end
TMPdls=MeanValdls;
TMPdls(isnan(TMPdls))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimgDLS]=sort(TMPdls);
DSconcat_OT_NPclassDLS=DSconcat_OT_NPclassDLS(SORTimgDLS,:);

DSconcat_OT_NPclassDMS=DSconcat_OTz(DMSsel,:);
for NN=1:size(DSconcat_OT_NPclassDMS,1)
    MeanValdms(NN,1)=mean(DSconcat_OT_NPclassDMS(NN,52:306),2);
end
TMPdms=MeanValdms;
TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
[~,SORTimgDMS]=sort(TMPdms);
DSconcat_OT_NPclassDMS=DSconcat_OT_NPclassDMS(SORTimgDMS,:);

% plot heatmap
subplot(6,4,[1 2 5 6])
imagesc(timeconcat,[1 size(DSconcat_OT_NPclassDLS,1)],DSconcat_OT_NPclassDLS,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([51 51],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
plot([102 102],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
plot([153 153],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
plot([204 204],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
plot([255 255],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
plot([306 306],[size(DSconcat_OT_NPclassDLS,1)+0.5 0.5],'k')
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
hold off

subplot(6,4,[9 10 13 14])
imagesc(timeconcat,[1 size(DSconcat_OT_NPclassDMS,1)],DSconcat_OT_NPclassDMS,Clim); %colorbar;axis tight;
colormap(jet);
hold on
plot([51 51],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
plot([102 102],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
plot([153 153],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
plot([204 204],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
plot([255 255],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
plot([306 306],[size(DSconcat_OT_NPclassDMS,1)+0.5 0.5],'k')
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
hold off

for i=1:max(eva.OptimalY)
    DSconcat_OT_classDLS=[]; SORTimgDLS=[];MaxValDLS=[];TMPdls=[];
    DSconcat_OT_classDMS=[]; SORTimgDMS=[];MaxValDMS=[];TMPdms=[];
    ind=find(eva.OptimalY==i);
    
    NPhasicNeurons_ID(:,1)=find(phasicID==1 & ~isnan(phasicID));
    NPhasicNeurons_ID(:,2)=Coord(NPhasicNeurons_ID(:,1),4);
    NPhasicNeurons_ID(:,3)=eva.OptimalY;
    
    DLSselectionNP=NPhasicNeurons_ID(:,2)==10 &  NPhasicNeurons_ID(:,3)==i;
    DMSselectionNP=NPhasicNeurons_ID(:,2)==20 &  NPhasicNeurons_ID(:,3)==i;
    
    
    % sorting DLS
    DSconcat_OT_classDLS=DSconcat_OT_NonPhasic(DLSselectionNP,:);
    for NN=1:size(DSconcat_OT_classDLS,1)
        MaxValDLS(NN,1)=mean(DSconcat_OT_classDLS(NN,52:306),2);
    end
    TMPdls=MaxValDLS;
    TMPdls(isnan(TMPdls))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimgDLS]=sort(TMPdls);
    DSconcat_OT_classDLS=DSconcat_OT_classDLS(SORTimgDLS,:);
    
    % sorting DMS
    DSconcat_OT_classDMS=DSconcat_OT_NonPhasic(DMSselectionNP,:);
    for NN=1:size(DSconcat_OT_classDMS,1)
        MaxValDMS(NN,1)=mean(DSconcat_OT_classDMS(NN,52:306),2);
    end
    TMPdms=MaxValDMS;
    TMPdms(isnan(TMPdms))=0; %To place the neurons with no onset/duration/peak at the top of the color-coded map
    [~,SORTimgDMS]=sort(TMPdms);
    DSconcat_OT_classDMS=DSconcat_OT_classDMS(SORTimgDMS,:);
    
    for j=1:2
        if j==1
            DSconcat_OT_class_DLSDMS= DSconcat_OT_classDLS;
        else
            DSconcat_OT_class_DLSDMS= DSconcat_OT_classDMS;
        end
        
        subplot(6,4,[3+(i-1)*12+(j-1)*4 4+(i-1)*12+(j-1)*4])
        imagesc(timeconcat,[1 size(DSconcat_OT_class_DLSDMS,1)],DSconcat_OT_class_DLSDMS,Clim);
        colormap(jet);
        hold on
        plot([51 51],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        plot([102 102],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        plot([153 153],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        plot([204 204],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        plot([255 255],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        plot([306 306],[size(DSconcat_OT_class_DLSDMS,1)+0.5 0.5],'k')
        set(gca,'XTick',[0:25.5:357]);
        set(gca,'xticklabel',Eventnames)
        hold off
        
        if i==1
            Yaxis=[-3 2];
        else
            Yaxis=[-1 15];
        end
        % plot average PSTH
        subplot(6,4,[11+(i-1)*12 12+(i-1)*12])
        DLSpsth=nanmean(DSconcat_OT_classDLS(:,timeconcat),1);
        DLSsem=nanste(DSconcat_OT_classDLS(:,timeconcat),1);
        DLSup=DLSpsth+DLSsem;
        DLSdown=DLSpsth-DLSsem;
        
        DMSpsth=nanmean(DSconcat_OT_classDMS(:,timeconcat),1);
        DMSsem=nanste(DSconcat_OT_classDMS(:,timeconcat),1);
        DMSup=DMSpsth+DMSsem;
        DMSdown=DMSpsth-DMSsem;
        hold on
        patch([timeconcat,timeconcat(end:-1:1)],[DLSup,DLSdown(end:-1:1)],myblue,'EdgeColor','none');alpha(0.5);
        g1=plot(timeconcat,DLSpsth,'Color',myblue,'linewidth',1.5);
        patch([timeconcat,timeconcat(end:-1:1)],[DMSup,DMSdown(end:-1:1)],'r','EdgeColor','none');alpha(0.5);
        g2=plot(timeconcat,DMSpsth,'r','linewidth',1.5);
        
        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([Xaxis2,Yaxis]);
        set(gca,'XTick',[0:25.5:357]);
        set(gca,'xticklabel',Eventnames)
        hold off
    end
end


%% proportion of neurons in each classes / plot stack bars / 2*3 classes with separation phasic / non phasic neurons

DSselection=[];

for j=1:2 % loop thru structure
    for k=1:6 % loop thru class
        DSselection(:,1)=NPhasicNeurons_ID(:,2)==BrainRegion(j) & NPhasicNeurons_ID(:,3)==k;
        DSselectionNumberNP(j,k)=sum(DSselection);
        DSselectionPropNP(j,k)=100*sum(DSselection)/sum(NPhasicNeurons_ID(:,2)==BrainRegion(j));
    end
end

subplot(6,4,[18 22])
X=categorical({'DLS','DMS'});
b1=bar(X,DSselectionPropNP,'stacked');
legend('INH','EXC');



%% make table for statistics
load('RatID_extendedtraining.mat');
 Class_OT(1:length(Coord(:,4)),1)=phasicID;
 Class_OT(PhasicNeurons_ID(:,1),2)=PhasicNeurons_ID(:,3);
 Class_OT(NPhasicNeurons_ID(:,1),2)=NPhasicNeurons_ID(:,3)+3;
 DSmeanz_OT=[];

 DSmeanz_OT = cat(2,DSmeanz_OT,Ev(7).Meanz(:,1)); % use meanZ instead of average of bin
  for k=1:5
      DSmeanz_OT = cat(2,DSmeanz_OT,Ev(k).MeanzPRE(:,1));
      DSmeanz_OT = cat(2,DSmeanz_OT,Ev(k).Meanz(:,1));
  end
 DSmeanz_OT = cat(2,DSmeanz_OT,Ev(6).MeanzPRE(:,1));


 table_OT=cat(2,RatID,Coord(:,4),Celltype(:,1),Class_OT,DSmeanz_OT, TRN);
 save('table_OT.mat','table_OT')

  nbevent=[1:12];
 
 selectionSTAT=table_OT(:,3)~=0 & table_OT(:,4)==1 & TRN(:,1)~=0;
 tableSTAT=array2table(table_OT(selectionSTAT,1:18),'VariableNames',{'rat','habit','region','celltype','phasicID','class','win1','win2','win3','win4','win5','win6','win7','win8','win9','win10','win11','win12'});

 rm = fitrm(tableSTAT,'win1-win12~region','WithinDesign',nbevent');
 ranovatbl_event = ranova(rm);   
 Btw_ranovatbl_event = anova(rm);
 
 %control for effect of habit
 rm = fitrm(tableSTAT,'win1-win12~region*habit','WithinDesign',nbevent');
 ranovatbl_event_habit = ranova(rm);   
 Btw_ranovatbl_event_habit = anova(rm);
 
 for i=1:max(table_OT(:,6))
    selectionSTAT=table_OT(:,3)~=0 & table_OT(:,4)==1 & table_OT(:,6)==i; 
    tableSTAT=array2table(table_OT(selectionSTAT,1:18),'VariableNames',{'rat','habit','region','celltype','phasicID','class','win1','win2','win3','win4','win5','win6','win7','win8','win9','win10','win11','win12'});

    rm = fitrm(tableSTAT,'win1-win12~region','WithinDesign',nbevent');
    stat(i).ranovatbl_event = ranova(rm);
    stat(i).Btw_ranovatbl_event = anova(rm);  
    
    rm = fitrm(tableSTAT,'win1-win12~region*habit','WithinDesign',nbevent');
    stat(i).ranovatbl_event_Habit = ranova(rm);
    stat(i).Btw_ranovatbl_event_Habit = anova(rm);   
 end


%% permutation test
allsamples=X_ratio;

for p=1:10000
    shuffsamples=c_ratio(randperm(length(c_ratio)));
    MeanCluster1(p,:)=mean(allsamples(shuffsamples==1,:),1);
    MeanCluster2(p,:)=mean(allsamples(shuffsamples==2,:),1);
    DistCluster(p,1) = pdist([MeanCluster1(p,:);MeanCluster2(p,:)],'euclidean');
end
MeanP=mean(allsamples(c_ratio==2,:),1);
MeanNP=mean(allsamples(c_ratio==1,:),1);
RealDist1=pdist([MeanP;MeanNP],'euclidean');
Classpval=(sum(RealDist1>DistCluster)+1)/(length(DistCluster)+1);

figure;
hist(DistCluster);
hold on;
plot([RealDist1 RealDist1],[0 4000]);
xlabel('Between-cluster distance')
ylabel('# iterations')
text(0.225,100,'Phasic vs Non-phasic classification','color','b','rotation',90);

%% Clustering on phasic and non phasic neurons 

rangeHighFreq=[2:0.1:8];
rangeLowFreq=[0.3:0.1:1.1];
nbClassif2=0;
for i=1:length(rangeHighFreq)
    for j=1:length(rangeLowFreq)
        
        f_low=rangeLowFreq(j);
        f_high=rangeHighFreq(i);
        
        %fourier analysis
        [X_all,f_all,P_all,Nt_all]=myfft(DSconcat_OTshort',0.01,0);
        total_power=sum(P_all);
        P_all_norm = bsxfun(@rdivide,P_all, total_power);
        
        a=find(f_all>f_low);b=find(f_all<f_high);
        lower_window_sum_all=sum(P_all_norm(1:a(1),:));
        upper_window_sum_all=sum(P_all_norm(a(1):b(end),:));
        
        ratio_all=lower_window_sum_all./upper_window_sum_all;
        
        %hierarchical clustering
        X_ratio=[[lower_window_sum_all]' [upper_window_sum_all]']; %#ok<NBRAK>
        eva_Ch_ratio = evalclusters(X_ratio,'linkage','CalinskiHarabasz','KList',[1:10]);
        OptimalK(i,j)=eva_Ch_ratio.OptimalK;
        c_ratio3=eva_Ch_ratio.OptimalY;
        if eva_Ch_ratio.OptimalK==2
            MeanP=mean(X_ratio(c_ratio3==2,:),1);
            MeanNP=mean(X_ratio(c_ratio3==1,:),1);
            RealDist(i,j)=pdist([MeanP;MeanNP],'euclidean');
            nbClassif2=nbClassif2+1;
            OptimalY(:,nbClassif2)=eva_Ch_ratio.OptimalY;
        end
    end
end

figure
[X,Y] = meshgrid(rangeLowFreq,rangeHighFreq);
pcolor(X,Y,OptimalK)
colorbar
xlabel('Lower frequency limit (Hz)')
ylabel('Higher frequency limit (Hz)')

corrOptimalK=OptimalK(:,1:8);
percentK2=100*sum(sum(corrOptimalK==2))/(size(corrOptimalK,1)*size(corrOptimalK,2));
percentK3=100*sum(sum(corrOptimalK==3))/(size(corrOptimalK,1)*size(corrOptimalK,2));

distancesDist=reshape(RealDist,size(RealDist,1)*size(RealDist,2),1);
figure;
hist(distancesDist(distancesDist(:,1)~=0,1));
hold on
hist(DistCluster);
plot([RealDist1 RealDist1],[0 4000]);
xlabel('Between-cluster distance')
ylabel('# iterations')
text(0.225,100,'Phasic vs Non-phasic classification','color','b','rotation',90);

A(1,1)=length(find(OptimalK==2));
A(1,2)=length(find(OptimalK==3));
A(1,3)=size(OptimalK,1)*size(OptimalK,2);
A(2,1)=100*A(1,1)./A(1,3);
A(2,2)=100*A(1,2)./A(1,3);
A(2,3)=100*(A(1,1)+A(1,2))./A(1,3);

%% look at 2 classes when OptimalK=2

for j=1:size(OptimalY,2)
    for i=1:2
        NbNeuronClass(j,i)=sum(OptimalY(:,j)==i);
    end
    NbNeuronClass(j,:)=sort(NbNeuronClass(j,:),2);
end

figure
histogram(NbNeuronClass(:,1))
hold on
histogram(NbNeuronClass(:,2))
hold off
xlabel('number of neurons in each class')
ylabel('# iterations')
legend('Non-Phasic','Phasic')