%% generate figure 8

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')
load('table_OT.mat')

%% --- --- Plotting color-coded and average PSTHs --- ---
BrainRegion=[10 20]; %10 for DLS, 20 for DMS
Xaxis=[0 357];Yaxis=[-2 5];
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'};
titlegraph={'Extended Training - sensitive', 'Extended Training - insensitive'};
time=-0.25:0.25:3.25;time2=1:1:357;

%% Extended training - Average psTH
%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

for i=1:2 %loop through habit/not habit neurons
DLSselection=Coord(:,4)==10 & Celltype(:,1)==1 & TRN(:,1)~=0 & RatID(:,2)==i-1;
DMSselection=Coord(:,4)==20 & Celltype(:,1)==1 & TRN(:,1)~=0 & RatID(:,2)==i-1;

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,:);


subplot(2,3,3+(i-1)*3)

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);
g1=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);
g2=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],[-2 5],'LineStyle',':','Color','k')
plot([102 102],[-2 5],'LineStyle',':','Color','k')
plot([153 153],[-2 5],'LineStyle',':','Color','k')
plot([204 204],[-2 5],'LineStyle',':','Color','k')
plot([255 255],[-2 5],'LineStyle',':','Color','k')
plot([306 306],[-2 5],'LineStyle',':','Color','k')
axis([Xaxis,Yaxis]);
set(gca,'XTick',[0:25.5:357]);
set(gca,'xticklabel',Eventnames)
traces=cat(2,g1,g2);
labels={'DLS','DMS'};
ylabel('z-score');
legend(traces,labels)
title(titlegraph(i))
end

DSconcat_OT = Ev(7).PSTHz(:,Ishow);
for i=1:6
    DSconcat_OT = cat(2,DSconcat_OT,Ev(i).PSTHz(:,Ishow));
end
LPwin=[52:280];
binLP=mean(DSconcat_OT(:,LPwin),2);

%% heatmap Extended training neurons
figure(2)

titleplot={'DLS sensitive', 'DMS sensitive', 'DLS insensitive', 'DMS insensitive'};
for i=1:max(table_OT(:,6))
    nbplottitle=0;
    for j=1:2 %loop thru habit profile
        for k=1:2 %loop thru region
            nbplottitle=nbplottitle+1;
            selection=Coord(:,4)==k*10 & RatID(:,2)==j-1 & table_OT(:,6)==i;
            DSconcat_Class_OT=DSconcat_OT(selection,:);
            
            %TMPdls=R.Class(DLSselection,1); TMPdms=R.Class(DMSselection,1);
            TMPds=binLP(selection,1);
            TMPds(isnan(TMPds))=0;
            [~,SORTimgDS]=sort(TMPds);[~,SORTimgDS]=sort(TMPds);
            DSconcat_Class_OT=DSconcat_Class_OT(SORTimgDS,:);
            
            subplot(4,5,1+(i-1)+10*(k-1)+5*(j-1))
            Clim=[-5 5];
            
            [nr,nc] = size(DSconcat_Class_OT);
            imagesc(time,[1 size(DSconcat_Class_OT,1)],DSconcat_Class_OT,Clim); %colorbar;axis tight;
            colormap(jet);
            hold on
            plot([0.25 0.25],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            plot([0.75 0.75],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            plot([1.25 1.25],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            plot([1.75 1.75],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            plot([2.25 2.25],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            plot([2.75 2.75],[size(DSconcat_Class_OT,1)+0.5 0.5],'k')
            set(gca,'XTick',[0:0.25:3.25]);
            set(gca,'xticklabel',Eventnames)
            ylabel('neurons');
            title(titleplot{nbplottitle})
            hold off
        end
    end
end

%% EARLY TRAININIG DATASET - average PSTH

clear Ev 
load ('Rearlytraining_light.mat');
load ('Celltype_earlyTraining.mat');
load('TRN_DT5_earlytrain.mat')
load('RatID_earlytraining.mat');
load('table_ACQ.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); 
    Ses(i).RatID(selectionIN,1)=RatID_ACQ(selectsession,1);  
end


%analysis restricted on putative MSN. nonTRN included
%concat with normalized data
nbplot=0;
figure(1)
titlegraph={'1st session - sensitive','last session - sensitive', '1st session - insensitive','last session - insensitive'};
for l=1:2 %loop thru habit profile
    for j=1:2 %loop thru 1st last session
        vars={'TMPdls','TMPdms','binSEQ','binLP','DLSselection','DMSselection','DSconcat_ACQ','DLSconcat_ACQ','DMSconcat_ACQ'};
        clear(vars{:});
        nbplot=nbplot+1;
        
        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 & Ses(k).RatID(:,1)==l-1;
        DMSselection=Ses(k).Coord(:,4)==20 & Ses(k).Celltype(:,1)==1 & sesTRN(k).TRN(:,1)~=0 & Ses(k).RatID(:,1)==l-1;
        
        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,:);
        
        
        
        %% Average OT PSTH
        traces=[];
        subplot(2,3,1+(j-1)+3*(l-1))
        
        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);
        g1=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);
        g2=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],[-2 5],'LineStyle',':','Color','k')
        plot([102 102],[-2 5],'LineStyle',':','Color','k')
        plot([153 153],[-2 5],'LineStyle',':','Color','k')
        plot([204 204],[-2 5],'LineStyle',':','Color','k')
        plot([255 255],[-2 5],'LineStyle',':','Color','k')
        plot([306 306],[-2 5],'LineStyle',':','Color','k')
        axis([Xaxis,Yaxis]);
        set(gca,'XTick',[0:25.5:357]);
        set(gca,'xticklabel',Eventnames)
        traces=cat(2,g1,g2);
        labels={'DLS','DMS'};
        ylabel('z-score');
        legend(traces,labels)
        title(titlegraph(nbplot))
        
    end
end

%% Early training data set - heatmaps


%Xaxis=[0 357];Yaxis=[-1 3];
Xaxis1=[-0.25 0.25];
Ishow=find(Tm>=Xaxis1(1) & Tm<=Xaxis1(2));
time=-0.25:0.25:3.25;
Eventnames={'LI';'';'LP1';'';'LP2';'';'LP3';'';'LP4';'';'LP5';'';'PE';''};


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); 
    Ses(i).RatID(selectionIN,1)=RatID_ACQ(selectsession,1); 
end


figure

for i=1:max(table_ACQ(:,5))
    nbtitle=0;
    for j=1:2 %loop thru habit profile
        for k=1:2 %loop thru region
            nbtitle=nbtitle+1;
            for l=4:13 %loop thru session
                DSconcat_ACQ=[];
                selection=Ses(l).Coord(:,4)==k*10 & Ses(l).Celltype(:,1)==1 & Ses(l).RatID(:,1)==j-1 & Ses(l).Class(:,1)==i;
                DSconcat_ACQ = Ses(l).Ev(7).PSTHz(selection,Ishow);
                for m=1:6
                    DSconcat_ACQ = cat(2,DSconcat_ACQ,Ses(l).Ev(m).PSTHz(selection,Ishow));
                end
                MeanDS.img(l-3,:)=nanmean(DSconcat_ACQ,1);
                if j==1
                    groupGD(k).region(l-3,i)=sum(selection);
                else
                    groupHAB(k).region(l-3,i)=sum(selection);
                end
            end
            
            subplot(4,5,1+(i-1)+10*(k-1)+5*(j-1))
            Clim=[-5 5];
            imagesc(time,[1 size(MeanDS.img,1)],MeanDS.img,Clim); %colorbar;axis tight;
            colormap(jet);
            hold on
            plot([0.25 0.25],[size(MeanDS.img,1)+0.5 0.5],'k')
            plot([0.75 0.75],[size(MeanDS.img,1)+0.5 0.5],'k')
            plot([1.25 1.25],[size(MeanDS.img,1)+0.5 0.5],'k')
            plot([1.75 1.75],[size(MeanDS.img,1)+0.5 0.5],'k')
            plot([2.25 2.25],[size(MeanDS.img,1)+0.5 0.5],'k')
            plot([2.75 2.75],[size(MeanDS.img,1)+0.5 0.5],'k')
            set(gca,'XTick',[0:0.25:3.25]);
            set(gca,'xticklabel',Eventnames)
            ylabel('neurons');
            title(titleplot(nbtitle));
            hold off
        end
    end
end