%% Plotting Heatmaps and average PSTH for sequence of event in DT5, FR5 and FS5
% edited on june 2019 to only include sucrose rats in DT5 data set
% do average first/last 3 session on PSTH instead of all sessions

clc;clear all; close all;

load('R_Comparison_2task.mat')
R2=R;
task_name={'DT5','FS5'};
task_Dataset={'R_DT5','R_FS5'};
BSIZE=0.01; %Do not change
Dura=[-25 20]; Tm=Dura(1):BSIZE:Dura(2);
thresholdIPI=[1,2,1];
nbLP=[2,3,2];

%% --- --- Plotting color-coded and average PSTHs --- ---
BrainRegion=[10 20]; %10 for DLS, 20 for DMS

% ----------- COLORS ---------
ColorsTask=[51/255 102/255 204/255,
            255/255 153/255 51/255];
        
Clim=[-2 2];%ColorMap
Xevent=[-0.25 0.25];
Yaxis=[-1 3];Xaxis=[1 357];
Ishow=find(Tm>=Xevent(1) & Tm<=Xevent(2));

% for stat
Ishow1(1,:)=find(Tm>=Xevent(1) & Tm<=0);
Ishow1(2,:)=find(Tm>=0 & Tm<=Xevent(2));
Combined_Table_allTask=[];allTask_ID=[];




%% ------------------------------------Main program ------------------------------------------------------------------------------------------------------------------------
figure;
for task=1:length(R)
    for event=1:3
        %generate concat file
        clear R Celltype SesCelltype 
        load(task_Dataset{task})
        DLS_activity_first3=[];DMS_activity_first3=[];DLS_activity_last3=[]; DMS_activity_last3=[];DLSconcat_ACQ_all=[];DMSconcat_ACQ_all=[];DLS_Sess_all=[];DMS_Sess_all=[];
        for k=1:7
            clear DLSselectionTRN DMSselectionTRN selectionTRN DLS_activity DMS_activity  
            for neuron=1:length(R2(task).Ses(k).IPI)
                ID_ShortIPI=R2(task).Ses(k).IPI{neuron,event}<thresholdIPI(event);
                if ~isempty(ID_ShortIPI)
                    IPI(neuron,1)=nanmean(R2(task).Ses(k).IPI{neuron,event}(ID_ShortIPI,1),1);
                else
                    IPI(neuron,1)=NaN;
                end
            end
            
            DLSselectionTRN=R.Ses(k).Coord(:,4)==BrainRegion(1) & R.Ses(k).TRN(:,1)~=0  & SesCelltype(k).Celltype(:,1)==1 & R2(task).Ses(k).NumberTrials(:,event)>5;
            DMSselectionTRN=R.Ses(k).Coord(:,4)==BrainRegion(2) & R.Ses(k).TRN(:,1)~=0  & SesCelltype(k).Celltype(:,1)==1 & R2(task).Ses(k).NumberTrials(:,event)>5;
            if ~isempty(R2(task).Ses(k).PSTHz{event})
                DLS_activity=R2(task).Ses(k).PSTHz{event}(DLSselectionTRN,Ishow);
                DMS_activity=R2(task).Ses(k).PSTHz{event}(DMSselectionTRN,Ishow);
            else
                DLS_activity=NaN(1,length(Ishow));
                DMS_activity=NaN(1,length(Ishow));
            end
            selectionTRN=R.Ses(k).Coord(:,4)~=0 & R.Ses(k).TRN(:,1)~=0  & SesCelltype(k).Celltype(:,1)==1 & R2(task).Ses(k).NumberTrials(:,event)>5;
            SelectionIPI=IPI(selectionTRN,1);
            ResponseRate{task}{k,event}=nbLP(event)./unique(SelectionIPI,'first');
            MeanRR{task}(k,event)=nanmean(ResponseRate{task}{k,event},1);
            semRR{task}(k,event)=nansem(ResponseRate{task}{k,event},1);
            
            if k<4 %first 3 DT5 sessions
                DLS_activity_first3=cat(1,DLS_activity_first3,DLS_activity);
                DMS_activity_first3=cat(1,DMS_activity_first3,DMS_activity);
            end
            if k>(length(R.Ses)-3)
                DLS_activity_last3=cat(1,DLS_activity_last3,DLS_activity);
                DMS_activity_last3=cat(1,DMS_activity_last3,DMS_activity);
            end
            DLSconcat_ACQ_all=cat(1,DLSconcat_ACQ_all,DLS_activity);
            DMSconcat_ACQ_all=cat(1,DMSconcat_ACQ_all,DMS_activity);
            DLS_Sess_all=cat(1,DLS_Sess_all,k*ones(size(DLS_activity,1),1));
            DMS_Sess_all=cat(1,DMS_Sess_all,k*ones(size(DMS_activity,1),1));
            time=[1:1:size(DLS_activity,2)];
            
            MeanDLS_ACQ{task}{event}(k,:)=nanmean(DLS_activity,1);
            MeanDMS_ACQ{task}{event}(k,:)=nanmean(DMS_activity,1);
            DLSsem1{task}{event}(k,:)=nansem(DLS_activity(:,time),1);
            DMSsem1{task}{event}(k,:)=nansem(DMS_activity(:,time),1);
            DLSnn{task}(k,event)=size(DLS_activity,1);
            DMSnn{task}(k,event)=size(DMS_activity,1);
        end
        TableAct_DLS{task,event}=DLSconcat_ACQ_all;
        TableAct_DMS{task,event}=DMSconcat_ACQ_all;
        TableSes_DLS{task,event}=DLS_Sess_all;
        TableSes_DMS{task,event}=DMS_Sess_all;
        
        
%         MeanDLS_ACQ_first3=nanmean(DLS_activity_first3,1);
%         MeanDMS_ACQ_first3=nanmean(DMS_activity_first3,1);
%         DLSsem1_first3=nansem(DLS_activity_first3,1);
%         DMSsem1_first3=nansem(DMS_activity_first3,1);
%         
%         MeanDLS_ACQ_last3=nanmean(DLS_activity_last3,1);
%         MeanDMS_ACQ_last3=nanmean(DMS_activity_last3,1);
%         DLSsem1_last3=nansem(DLS_activity_last3,1);
%         DMSsem1_last3=nansem(DMS_activity_last3,1);
        
        MeanDLS_ACQ_all=nanmean(DLSconcat_ACQ_all,1);
        MeanDMS_ACQ_all=nanmean(DMSconcat_ACQ_all,1);
        DLSsem1_all=nansem(DLSconcat_ACQ_all,1);
        DMSsem1_all=nansem(DMSconcat_ACQ_all,1);
        

        %% Heatmaps
        subplot(4,6,6+event+6*(task-1))
        imagesc(time,[1 size(MeanDLS_ACQ{task}{event},1)],MeanDLS_ACQ{task}{event},Clim); %colorbar;axis tight;
        colormap(parula);
        hold on
        plot([0 0],[size(MeanDLS_ACQ{task}{event},1)+0.5 0.5],'k')
        hold off
        set(gca,'XTick',0:25.5:size(MeanDLS_ACQ{task}{event},2));
        set(gca,'xticklabel',{'-0.25','0','0.25'})
        set(gca,'YTick',1:size(MeanDLS_ACQ{task}{event},1));
        set(gca,'yticklabel',DLSnn{task}(:,event))
        ylabel('sessions');
        
        subplot(4,6,9+event+6*(task-1))
        imagesc(time,[1 size(MeanDMS_ACQ{task}{event},1)],MeanDMS_ACQ{task}{event},Clim); %colorbar;axis tight;
        colormap(parula);
        hold on
        plot([0 0],[size(MeanDMS_ACQ{task}{event},1)+0.5 0.5],'k')
        set(gca,'XTick',0:25.5:size(MeanDMS_ACQ{task}{event},2));
        set(gca,'xticklabel',{'-0.25','0','0.25'})
        set(gca,'YTick',1:size(MeanDMS_ACQ{task}{event},1));
        set(gca,'yticklabel',DMSnn{task}(:,event))
        hold off
        
        subplot(4,6,18+event)
        DLSup=MeanDLS_ACQ_all+DLSsem1_all;
        DLSdown=MeanDLS_ACQ_all-DLSsem1_all;
        hold on;
        patch([time,time(end:-1:1)],[DLSup,DLSdown(end:-1:1)],ColorsTask(task,:),'EdgeColor','none');alpha(0.5);
        g1=plot(time,MeanDLS_ACQ_all,'Color',ColorsTask(task,:),'linewidth',1.5);
        plot([0 50],[0 0],'k','LineStyle',':')
        plot([25 25],[-1 2.5],'k','LineStyle',':')
        axis([0 51 -1 2.5]);
        set(gca,'XTick',0:25.5:size(MeanDMS_ACQ{task}{event},2));
        set(gca,'xticklabel',{'-0.25','0','0.25'})
        ylabel('z-score');
        ylim([-1 2.5]);
        
        subplot(4,6,21+event)
        DMSup=MeanDMS_ACQ_all+DMSsem1_all;
        DMSdown=MeanDMS_ACQ_all-DMSsem1_all;
        hold on;
        patch([time,time(end:-1:1)],[DMSup,DMSdown(end:-1:1)],ColorsTask(task,:),'EdgeColor','none');alpha(0.5);
        g2=plot(time,MeanDMS_ACQ_all,'Color',ColorsTask(task,:),'linewidth',1.5);
        plot([0 50],[0 0],'k','LineStyle',':')
        plot([25 25],[-1 2.5],'k','LineStyle',':')
        axis([0 51 -1 2.5]);
        set(gca,'XTick',0:25.5:size(MeanDMS_ACQ{task}{event},2));
        set(gca,'xticklabel',{'-0.25','0','0.25'})
        ylabel('z-score');        
        hold off      
    end
end

for event=1:3
    for task=1:2
        subplot(4,6,3+event)
        hold on
        er = errorbar(1:7,MeanRR{task}(:,event),semRR{task}(:,event),'-o','MarkerSize',4,'color',ColorsTask(task,:),'MarkerFaceColor',ColorsTask(task,:),'CapSize',0,'LineWidth',1);
        ylabel('response rate (Hz)');
        xlabel('sessions')
        axis([0 8 0 10])
    end
end

%% statistics
filename = 'ResultStat_Activity_Fig5.xlsx'; 
for event=1:3
    ActDLS_BothTask=cat(1,TableAct_DLS{:,event});
    ActDMS_BothTask=cat(1,TableAct_DMS{:,event});
    TableActDLS_BothTask=[nanmean(ActDLS_BothTask(:,1:25),2) nanmean(ActDLS_BothTask(:,26:51),2)];
    TableActDMS_BothTask=[nanmean(ActDMS_BothTask(:,1:25),2) nanmean(ActDMS_BothTask(:,26:51),2)];
    Task_ID_DLS=cat(1,zeros(size(TableAct_DLS{1,event},1),1),ones(size(TableAct_DLS{2,event},1),1));
    Task_ID_DMS=cat(1,zeros(size(TableAct_DMS{1,event},1),1),ones(size(TableAct_DMS{2,event},1),1));
    SessionsDLS_BothTask=cat(1,TableSes_DLS{:,event});
    SessionsDMS_BothTask=cat(1,TableSes_DMS{:,event});
    
    CatDLS{event}=cat(2,Task_ID_DLS,SessionsDLS_BothTask,TableActDLS_BothTask);
    CatDMS{event}=cat(2,Task_ID_DMS,SessionsDMS_BothTask,TableActDMS_BothTask);
    
    nbwin=[1:2];
    %DLS first
    tableSTAT=array2table(CatDLS{event},'VariableNames',{'task','session','win1','win2'});
    tableSTAT.session=categorical(tableSTAT.session);
    tableSTAT.task=categorical(tableSTAT.task);
    
    rm_reg = fitrm(tableSTAT,'win1-win2~task*session','WithinDesign',nbwin');
    stat(event).condition={'rmANOVA DLS'};
    stat(event).anovatbl_event = ranova(rm_reg);
    stat(event).Btw_ranovatbl_event = anova(rm_reg);
    
    writetable(stat(event).anovatbl_event,filename,'Sheet',event,'Range','A1')
    writetable(stat(event).Btw_ranovatbl_event,filename,'Sheet',event,'Range','J1')
    
    %then test DMS
    tableSTAT=array2table(CatDMS{event},'VariableNames',{'task','session','win1','win2'});
    tableSTAT.session=categorical(tableSTAT.session);
    tableSTAT.task=categorical(tableSTAT.task);
    
    rm_reg = fitrm(tableSTAT,'win1-win2~task*session','WithinDesign',nbwin');
    stat(event+3).condition={'rmANOVA DMS'};
    stat(event+3).anovatbl_event = ranova(rm_reg);
    stat(event+3).Btw_ranovatbl_event = anova(rm_reg);
    
    writetable(stat(event+3).anovatbl_event,filename,'Sheet',event+3,'Range','A1')
    writetable(stat(event+3).Btw_ranovatbl_event,filename,'Sheet',event+3,'Range','J1')
end

%Stat RR
SesID=[];TableRR=[];taskID=[];
for event=1:3
    for task=1:2
        SesID_sh=[];
        for k=1:7
            SesID=cat(1,SesID,k*ones(length(ResponseRate{task}{k,event}),1));
            TableRR=cat(1,TableRR,ResponseRate{task}{k,event});
            SesID_sh=cat(1,SesID_sh,k*ones(length(ResponseRate{task}{k,event}),1));
        end
        taskID=cat(1,taskID,task*ones(length(SesID_sh),1));
    end
    TableStat_RR{event}=cat(2,taskID,SesID,TableRR);

    stat(6+event).condition={'ANOVAN RR'};
    [~,stat(6+event).anovatbl_event,~]=anovan(TableStat_RR{event}(:,3),{TableStat_RR{event}(:,1),TableStat_RR{event}(:,2)},'varnames',{'task','session'},'display','off','model','full');
    xlswrite(filename,stat(6+event).anovatbl_event,event,'R1')
end