%% generate figure 7A. Decoding brain region ID without PCA

% code from David Ottenheimer 04/19/2018

%get the data
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).Meanz(DLSselection,1));
    DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(7).Meanz(DMSselection,1));
    for j=1:5
        DLSconcat_ACQ = cat(2,DLSconcat_ACQ,Ses(i).Ev(j).MeanzPRE(DLSselection,1));
        DLSconcat_ACQ = cat(2,DLSconcat_ACQ,Ses(i).Ev(j).Meanz(DLSselection,1));
        DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(j).MeanzPRE(DMSselection,1));
        DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(j).Meanz(DMSselection,1));
    end
    DLSconcat_ACQ = cat(2,DLSconcat_ACQ,Ses(i).Ev(6).MeanzPRE(DLSselection,1));
    DMSconcat_ACQ = cat(2,DMSconcat_ACQ,Ses(i).Ev(6).MeanzPRE(DMSselection,1));

    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).Meanz(DLSselection,1);
DMSconcat_OT = Ev(7).Meanz(DMSselection,1);
for i=1:5
    DLSconcat_OT = cat(2,DLSconcat_OT,Ev(i).MeanzPRE(DLSselection,1));
    DLSconcat_OT = cat(2,DLSconcat_OT,Ev(i).Meanz(DLSselection,1));
    DMSconcat_OT = cat(2,DMSconcat_OT,Ev(i).MeanzPRE(DMSselection,1));
    DMSconcat_OT = cat(2,DMSconcat_OT,Ev(i).Meanz(DMSselection,1));
end
DLSconcat_OT = cat(2,DLSconcat_OT,Ev(6).MeanzPRE(DLSselection,1));
DMSconcat_OT = cat(2,DMSconcat_OT,Ev(6).MeanzPRE(DMSselection,1));

MeanzDecodeOT.DLS=DLSconcat_OT;
MeanzDecodeOT.DMS=DMSconcat_OT;

save('MeanzDecodeACQ','MeanzDecodeOT')



%% ACQ PCcomponents Comparison sessions. 
%PCA first
%first do PCA on an equal number of neurons in each region
nbsession=0;
for k=1:10
    clear rep
    countRep=0;
    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);
        
        rep(l).concat=cat(1,PSTHzDecode.DLS(NSel,:),PSTHzDecode.DMS(VSel,:));
    end
    
    %% Decoding
    clear Partitions SVMModel prediction actual correct DecodeRegion
    folds = 10;%PCAneurons; %number of times cross-validated
    shuffs = 1; %number of times shuffled
    DiscrimType= 'linear';
    
    %setup the event identity matrix for decoding
    DecodeRegion(1:PCAneurons,1)=1;
    DecodeRegion((PCAneurons+1):(PCAneurons*2),1)=2;
    for l=1:50
        
            countRep=countRep+1;
            clear CVacc CVaccSh
            %Setup spikes matrix for decoding
            DecodeMeanZ=rep(l).concat;
            
            %normal model
            for r = 1:folds
                
                Partitions = cvpartition(DecodeRegion,'KFold',folds);
                SVMModel = fitcdiscr(DecodeMeanZ(Partitions.training(r),:),DecodeRegion(Partitions.training(r)),'DiscrimType',DiscrimType);
                prediction = predict(SVMModel,DecodeMeanZ(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(DecodeMeanZ(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)),'DiscrimType',DiscrimType);
                    predictionSh = predict(SVMModelSh,DecodeMeanZ(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{k,1}(countRep,1)=nanmean(CVacc);
            PCADecodeAccShAA{k,1}(countRep,1)=nanmean(AccShuff);
            fprintf(['RepSelection #' num2str(countRep) '\n']);
        end
    
end


%% plotting
%colors
inh=[0 0 0.6];
exc=[0.8 0 0];
NumSession=[1 2 3 4 5 6 7 8 9 10];


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


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

for k=1:5
    countRep=0;
    clear rep
    
    PSTHzDecode.DLS=MeanzDecodeOT.DLS;
    PSTHzDecode.DMS=MeanzDecodeOT.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
    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);
    
    rep(l).concatOT=cat(1,PSTHzDecode.DLS(NSel,:),PSTHzDecode.DMS(VSel,:));
    end
    
    %% Decoding
    clear Partitions SVMModel prediction actual correct DecodeRegion
    folds = 10;%PCAneurons; %number of times cross-validated
    shuffs = 1; %number of times shuffled
    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;
    for l=1:50
   
        countRep=countRep+1;

        clear CVacc CVaccSh
        %Setup spikes matrix for decoding
        DecodeMeanZ=rep(l).concatOT;
        
        %normal model
        for r = 1:folds
            
            Partitions = cvpartition(DecodeRegion,'KFold',folds);
            SVMModel = fitcdiscr(DecodeMeanZ(Partitions.training(r),:),DecodeRegion(Partitions.training(r)),'DiscrimType',DiscrimType);
            prediction = predict(SVMModel,DecodeMeanZ(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(DecodeMeanZ(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)),'DiscrimType',DiscrimType);
                predictionSh = predict(SVMModelSh,DecodeMeanZ(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{1,k}(countRep,1)=nanmean(CVacc);
        PCADecodeAccShOT{1,k}(countRep,1)=nanmean(AccShuff);
        
        fprintf(['Rep #' num2str(countRep) '\n']);
    end
   
end


%% plotting
%colors
inh=[0 0 0.6];
exc=[0.8 0 0];

Shuffle=[];
for k=1:5
    subplot(1,2,2);
    hold on;
    errorbar(nbneurons(k),nanmean(PCADecodeAccOT{1,k}),nanstd(PCADecodeAccOT{1,k}),'o','color',exc(1,:));
    errorbar(nbneurons(k),nanmean(PCADecodeAccShOT{1,k}),nanstd(PCADecodeAccShOT{1,k}),'o','color',inh(1,:));
end

axis([0.5 300 0.35 0.85]);
title('Ext training: Decoding DMS vs DLS');
xlabel('Ext Training - ensemble size');
ylabel('Accuracy');

%save('Decode_withoutPCA.mat','PCADecodeAccAA','PCADecodeAccOT','PCADecodeAccShAA','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{1,i},PCADecodeAccShOT{1,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{1,1});
[p_ACQ_OT,tbl_ACQ_OT]=anova1(tableAccACQ_OT,Between_factor,'off');

%%

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
for i=1:length(PCADecodeAccOT)
    allsamples=cat(1,PCADecodeAccOT{1,i},PCADecodeAccShOT{1,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);
end