%% Generate figure 06 and S07
% decoding using PCA; Up left, compare 3 early training sessions across PCs 
% Up right, compare different ensemble size of extended training across PCs 


% code from David Ottenheimer 04/19/2018
% decode region
% first I'll try just classifying individual neurons as DMS or DLS
%if ~exist('PSTHzDecode') %if you haven't already collected all the psth data for the decoding

clear all
close all

load('Rearlytraining_light.mat')
Xaxis1=[-0.25 0.25];
Ishow=find(Tm>=Xaxis1(1) & Tm<=Xaxis1(2));


for i=4:13 %loop thru ACQ session

   DLSselection=Ses(i).Coord(:,4)==10 & Ses(i).Celltype(:,1)==1;
   DMSselection=Ses(i).Coord(:,4)==20 & Ses(i).Celltype(:,1)==1;
    
    DLSconcat_ACQ=[];DMSconcat_ACQ=[];
    DLSconcat_ACQ = cat(2,DLSconcat_ACQ,Ses(i).Ev(7).PSTHz(DLSselection,Ishow));
    DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(7).PSTHz(DMSselection,Ishow));
    for j=1:6
        DLSconcat_ACQ = cat(2,DLSconcat_ACQ,Ses(i).Ev(j).PSTHz(DLSselection,Ishow));
        DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(j).PSTHz(DMSselection,Ishow));
    end

    PSTHzDecodeACQ(i-3).DLS=DLSconcat_ACQ(~isnan(mean(DLSconcat_ACQ,2)),:);
    PSTHzDecodeACQ(i-3).DMS=DMSconcat_ACQ(~isnan(mean(DMSconcat_ACQ,2)),:);

end

clear Ses
load('Rextendedtraining_light.mat');
load('Celltype_extendedtraining.mat');

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

DLSconcat_OT = Ev(7).PSTHz(DLSselection,Ishow);
DMSconcat_OT = Ev(7).PSTHz(DMSselection,Ishow);
for i=1:6
    DLSconcat_OT = cat(2,DLSconcat_OT,Ev(i).PSTHz(DLSselection,Ishow));
    DMSconcat_OT = cat(2,DMSconcat_OT,Ev(i).PSTHz(DMSselection,Ishow));
end
PSTHzDecodeOT.DLS=DLSconcat_OT;
PSTHzDecodeOT.DMS=DMSconcat_OT;

save('PSTHzDecodeACQ','PSTHzDecodeOT')



%% ACQ PCcomponents Comparison sessions. 
%PCA first
%first do PCA on an equal number of neurons in each region
nbsession=0;
for k=1:3:10
    countRep=0;
    nbsession=nbsession+1;
    
    PSTHzDecode.DLS=PSTHzDecodeACQ(k).DLS;
    PSTHzDecode.DMS=PSTHzDecodeACQ(k).DMS;
    
    if length(PSTHzDecode.DLS(:,1))<length(PSTHzDecode.DMS(:,1))
        PCAneurons=length(PSTHzDecode.DLS(:,1));
    else
        PCAneurons=length(PSTHzDecode.DMS(:,1));
    end
    
    for l=1:50 % loop added to account for selection bias in VSel and NSel
        %pick which neurons to use
        SetupDMSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DMS(:,1))-PCAneurons,1));
        VSel=(SetupDMSel(randperm(length(SetupDMSel)))==1);
        
        %pick which neurons to use
        SetupDLSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DLS(:,1))-PCAneurons,1));
        NSel=(SetupDLSel(randperm(length(SetupDLSel)))==1);
        
        PCAneuronsACQ(nbsession,1)=PCAneurons;
        
        concat=cat(1,PSTHzDecode.DLS(NSel,:),PSTHzDecode.DMS(VSel,:))';
        [coeff,score,~,~,explained] = pca(concat);
        
        
        % Decoding
        clear Partitions SVMModel prediction actual correct
        folds = 10;%PCAneurons; %number of times cross-validated
        shuffs = 1; %number of times shuffled
        NumCoeffs = [1 2 3 4 5 6 7 8 9 10]; %number of PCs used in decoding
        repetitions=20; %how many times to run the analysis
        DiscrimType= 'linear';
        
        %setup the event identity matrix for decoding
        DecodeRegion(1:PCAneurons,1)=1;
        DecodeRegion((PCAneurons+1):(PCAneurons*2),1)=2;
        
        
        countRep=countRep+1;
        
        for i=1:length(NumCoeffs)
            clear CVacc CVaccSh
            %Setup spikes matrix for decoding
            DecodePCs=coeff(:,1:NumCoeffs(i));
            %ADD IN PCA HERE???
            
            %normal model
            for r = 1:folds
                
                Partitions = cvpartition(DecodeRegion,'KFold',folds);
                SVMModel = fitcdiscr(DecodePCs(Partitions.training(r),:),DecodeRegion(Partitions.training(r)),'DiscrimType',DiscrimType);
                prediction = predict(SVMModel,DecodePCs(Partitions.test(r),:));
                actual = DecodeRegion(Partitions.test(r));
                correct = prediction - actual;
                CVacc(r,1) = sum(correct==0) / length(correct);
                
            end
            
            %shuffled model
            for q=1:shuffs
                DecodeRsSh=DecodeRegion(randperm(length(DecodeRegion)));
                PartitionsSh = cvpartition(DecodeRsSh,'KFold',folds);
                for s = 1:folds
                    SVMModelSh = fitcdiscr(DecodePCs(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)),'DiscrimType',DiscrimType);
                    predictionSh = predict(SVMModelSh,DecodePCs(PartitionsSh.test(s),:));
                    actualSh = DecodeRsSh(PartitionsSh.test(s));
                    correctSh = predictionSh - actualSh;
                    CVaccSh(s,1) = sum(correctSh==0) / length(correctSh);
                end
                AccShuff(q,1) = nanmean(CVaccSh);
            end
            PCADecodeAccACQ{i,nbsession}(countRep,1)=nanmean(CVacc);
            PCADecodeAccShACQ{i,nbsession}(countRep,1)=nanmean(AccShuff);
        end
        fprintf(['RepACQ #' num2str(countRep) '\n']);
    end  
end

%plotting
%colors
inh=[0 0.3333 1; 0 0.5 1; 0 0.67 1; 0 0.8353 1];
exc=[1 0.1647 0; 0.9020 0.4510 0; 0.9020 0.6 0; 1 0.8353 0];

%Y = prctile(X,p)

for i=1:length(NumCoeffs)
    for k=1:4
        subplot(2,8,[1 2 3 4]);
        hold on;
        errorbar(NumCoeffs(i),nanmean(PCADecodeAccACQ{i,k}),nanstd(PCADecodeAccACQ{i,k}),'o','color',exc(k,:));
        errorbar(NumCoeffs(i),nanmean(PCADecodeAccShACQ{i,4}),nanstd(PCADecodeAccShACQ{i,4}),'o','color',inh(1,:));
    end
end

axis([0 NumCoeffs(end)+1 0.30 0.85]);
title('Early training: Decoding DMS vs DLS');
xlabel('Number of PCs');
ylabel('Accuracy');

%% Overtraining as a function of PC and number of neurons included

for k=1:5
    countRep=0;
    
    PSTHzDecode.DLS=PSTHzDecodeOT.DLS;
    PSTHzDecode.DMS=PSTHzDecodeOT.DMS;
    
    nbneurons=[30 60 100 200 300];
    PCAneurons=nbneurons(k);
    
    for l=1:50 % loop added to account for selection bias in VSel and NSel
        %pick which neurons to use
        SetupDMSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DMS(:,1))-PCAneurons,1));
        VSel=(SetupDMSel(randperm(length(SetupDMSel)))==1);
        
        %pick which neurons to use
        SetupDLSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DLS(:,1))-PCAneurons,1));
        NSel=(SetupDLSel(randperm(length(SetupDLSel)))==1);
        
        concat=cat(1,PSTHzDecode.DLS(NSel,:),PSTHzDecode.DMS(VSel,:))';
        [coeff,score,~,~,explained] = pca(concat);
        
        %% Decoding
        clear Partitions SVMModel prediction actual correct DecodeRegion
        folds = 10;%PCAneurons; %number of times cross-validated
        shuffs = 1; %number of times shuffled
        NumCoeffs = [1 2 3 4 5 6 7 8 9 10]; %number of PCs used in decoding
        DiscrimType= 'linear';
        
        %setup the event identity matrix for decoding
        DecodeRegion(1:PCAneurons,1)=1;
        DecodeRegion((PCAneurons+1):(PCAneurons*2),1)=2;
        
        countRep=countRep+1;
        
        for i=1:length(NumCoeffs)
            
            clear CVacc CVaccSh
            %Setup spikes matrix for decoding
            DecodePCs=coeff(:,1:NumCoeffs(i));
            %ADD IN PCA HERE???
            
            %normal model
            for r = 1:folds
                
                Partitions = cvpartition(DecodeRegion,'KFold',folds);
                SVMModel = fitcdiscr(DecodePCs(Partitions.training(r),:),DecodeRegion(Partitions.training(r)),'DiscrimType',DiscrimType);
                prediction = predict(SVMModel,DecodePCs(Partitions.test(r),:));
                actual = DecodeRegion(Partitions.test(r));
                correct = prediction - actual;
                CVacc(r,1) = sum(correct==0) / length(correct);
                
            end
            
            %shuffled model
            for q=1:shuffs
                DecodeRsSh=DecodeRegion(randperm(length(DecodeRegion)));
                PartitionsSh = cvpartition(DecodeRsSh,'KFold',folds);
                for s = 1:folds
                    SVMModelSh = fitcdiscr(DecodePCs(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)),'DiscrimType',DiscrimType);
                    predictionSh = predict(SVMModelSh,DecodePCs(PartitionsSh.test(s),:));
                    actualSh = DecodeRsSh(PartitionsSh.test(s));
                    correctSh = predictionSh - actualSh;
                    CVaccSh(s,1) = sum(correctSh==0) / length(correctSh);
                end
                AccShuff(q,1) = nanmean(CVaccSh);
            end
            PCADecodeAccOT{i,k}(countRep,1)=nanmean(CVacc);
            PCADecodeAccShOT{i,k}(countRep,1)=nanmean(AccShuff);
        end
        fprintf(['RepOT #' num2str(countRep) '\n']);
    end
end


%% plotting
%colors
inh=[0 0 0.6; 0 0.3333 1; 0 0.5 1; 0 0.67 1; 0 0.8353 1];
exc=[0.8 0 0; 1 0.1647 0; 0.9020 0.4510 0; 0.9020 0.6 0; 1 0.8353 0];


for i=1:length(NumCoeffs)
    for k=1:5
        subplot(2,8,[5 6 7 8]);
        hold on;
        errorbar(NumCoeffs(i),nanmean(PCADecodeAccOT{i,k}),nanstd(PCADecodeAccOT{i,k}),'o','color',exc(k,:));
        errorbar(NumCoeffs(i),nanmean(PCADecodeAccShOT{i,5}),nanstd(PCADecodeAccShOT{i,5}),'o','color',inh(1,:));
    end
end

axis([0 NumCoeffs(end)+1 0.30 0.85]);
title('Ext Training: Decoding DMS vs DLS');
xlabel('Number of PCs');
ylabel('Accuracy');



%% Decoding DMS vs DLS across session, based on weight of 3 first PCs
%PCA first
for i=1:10
    countRep=0;
    
    PSTHzDecode.DLS=PSTHzDecodeACQ(i).DLS;
    PSTHzDecode.DMS=PSTHzDecodeACQ(i).DMS;
    
    if length(PSTHzDecode.DLS(:,1))<length(PSTHzDecode.DMS(:,1))
        PCAneurons=length(PSTHzDecode.DLS(:,1));
    else
        PCAneurons=length(PSTHzDecode.DMS(:,1));
    end
    
    for l=1:50
        SetupDMSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DMS(:,1))-PCAneurons,1));
        VSel=(SetupDMSel(randperm(length(SetupDMSel)))==1);
        
        %pick which neurons to use
        SetupDLSel=cat(1,ones(PCAneurons,1),zeros(length(PSTHzDecode.DLS(:,1))-PCAneurons,1));
        NSel=(SetupDLSel(randperm(length(SetupDLSel)))==1);
        
        concat=cat(1,PSTHzDecode.DLS(NSel,:),PSTHzDecode.DMS(VSel,:))';
        [coeff,score,~,~,explained] = pca(concat);
        
        %% Decoding
        clear Partitions SVMModel prediction actual correct DecodeRegion
        folds = 10;%PCAneurons; %number of times cross-validated
        shuffs = 1; %number of times shuffled
        NumCoeffs = 3; %number of PCs used in decoding
        DiscrimType= 'linear';
        
        %setup the event identity matrix for decoding
        DecodeRegion(1:PCAneurons,1)=1;
        DecodeRegion((PCAneurons+1):(PCAneurons*2),1)=2;
        
        countRep=countRep+1;
        
        clear CVacc CVaccSh
        %Setup spikes matrix for decoding
        DecodePCs=coeff(:,1:NumCoeffs);
        %ADD IN PCA HERE???
        
        %normal model
        for r = 1:folds
            
            Partitions = cvpartition(DecodeRegion,'KFold',folds);
            SVMModel = fitcdiscr(DecodePCs(Partitions.training(r),:),DecodeRegion(Partitions.training(r)),'DiscrimType',DiscrimType);
            prediction = predict(SVMModel,DecodePCs(Partitions.test(r),:));
            actual = DecodeRegion(Partitions.test(r));
            correct = prediction - actual;
            CVacc(r,1) = sum(correct==0) / length(correct);
            
        end
        
        %shuffled model
        for q=1:shuffs
            DecodeRsSh=DecodeRegion(randperm(length(DecodeRegion)));
            PartitionsSh = cvpartition(DecodeRsSh,'KFold',folds);
            for s = 1:folds
                SVMModelSh = fitcdiscr(DecodePCs(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)),'DiscrimType',DiscrimType);
                predictionSh = predict(SVMModelSh,DecodePCs(PartitionsSh.test(s),:));
                actualSh = DecodeRsSh(PartitionsSh.test(s));
                correctSh = predictionSh - actualSh;
                CVaccSh(s,1) = sum(correctSh==0) / length(correctSh);
            end
            AccShuff(q,1) = nanmean(CVaccSh);
        end
        PCADecodeAccAA{i,1}(countRep,1)=nanmean(CVacc);
        PCADecodeAccShAA{i,1}(countRep,1)=nanmean(AccShuff);
        fprintf(['RepACQses #' num2str(countRep) '\n']);
    end
end


%% plotting
%colors
inh=[0 0 0.6; 0 0.3333 1; 0 0.5 1; 0 0.67 1; 0 0.8353 1];
exc=[0.8 0 0; 1 0.1647 0; 0.9020 0.4510 0; 0.9020 0.6 0; 1 0.8353 0];
NumSession=[1 2 3 4 5 6 7 8 9 10];

for i=1:10
    subplot(2,8,[9 10 11 12]);
    hold on;
    errorbar(NumSession(i),nanmean(PCADecodeAccAA{i,1}),nanstd(PCADecodeAccAA{i,1}),'o','color',exc(1,:));
    errorbar(NumSession(i),nanmean(PCADecodeAccShAA{i,1}),nanstd(PCADecodeAccShAA{i,1}),'o','color',inh(1,:));
end
axis([0 10+1 0.35 0.85]);
title('Decoding of region across session');
xlabel('Sessions');
ylabel('Accuracy');


for k=1:5
subplot(2,8,[13 14 15]);
hold on;
errorbar(nbneurons(k),nanmean(PCADecodeAccOT{3,k}),nanstd(PCADecodeAccOT{3,k}),'o','color',exc(1,:));
errorbar(nbneurons(k),nanmean(PCADecodeAccShOT{3,k}),nanstd(PCADecodeAccShOT{3,k}),'o','color',inh(1,:));
end
axis([0 300 0.35 0.85]);
xlabel('OT');
ylabel('Accuracy');

save('PCADecode.mat','PCADecodeAccAA','PCADecodeAccACQ','PCADecodeAccOT','PCADecodeAccShAA','PCADecodeAccShACQ','PCADecodeAccShOT');

%% stat
for i=1:10
   tableAcc(:,i)=cat(1,PCADecodeAccAA{i,1},PCADecodeAccShAA{i,1});
   Between_factor=cat(1,zeros(50,1),ones(50,1));
   [p,tbl]=anova1(tableAcc(:,i),Between_factor,'off');
   p_value(i)=p;
   stat(i).t=tbl;
end

for i=1:5
    tableAccOT(:,i)=cat(1,PCADecodeAccOT{3,i},PCADecodeAccShOT{3,i});
    [p,tbl]=anova1(tableAccOT(:,i),Between_factor,'off');
    pOT_value(i)=p;
    statOT(i).t=tbl;
end

tableAccACQ_OT=cat(1,PCADecodeAccAA{10,1},PCADecodeAccOT{3,1});
[p_ACQ_OT,tbl_ACQ_OT]=anova1(tableAccACQ_OT,Between_factor,'off');

%% Permutation test
for i=1:length(PCADecodeAccAA)
    allsamples=cat(1,PCADecodeAccAA{i,1},PCADecodeAccShAA{i,1});
    for p=1:10000
        shuffsamples=allsamples(randperm(length(allsamples)));
        shuffdiff(p,1)=abs(mean(shuffsamples(1:length(allsamples)/2))-mean(shuffsamples(length(allsamples)/2+1:end)));
    end
    diff=abs(nanmean(PCADecodeAccAA{i,1})-nanmean(PCADecodeAccShAA{i,1}));
    trainingpval(i,1)=(sum(shuffdiff>diff)+1)/(length(shuffdiff)+1);
end

%% overtraining
figure;

for i=1:size(PCADecodeAccOT,2)
    allsamples=cat(1,PCADecodeAccOT{3,i},PCADecodeAccShOT{3,i});
    for p=1:10000
        shuffsamples=allsamples(randperm(length(allsamples)));
        shuffdiff(p,1)=abs(mean(shuffsamples(1:length(allsamples)/2))-mean(shuffsamples(length(allsamples)/2+1:end)));
    end
    diff=abs(nanmean(PCADecodeAccOT{1,i})-nanmean(PCADecodeAccShOT{1,i}));
    OTpval(1,i)=(sum(shuffdiff>diff)+1)/(length(shuffdiff)+1);
    
    subplot(5,1,i)
    hist(shuffdiff);
    hold on;
    plot([diff diff],[0 4000]);
end