%decodes trial identity from spiking across bins for neurons pooled across sessions
%also does this for only selective and non-selective neurons

%need to run A, B, and D first
clear all;
load ('R_2R.mat');
load ('RAW.mat');

%NAc is 1, VP is 2
region={'NA';'VP'};

%decoding parameters
NumNeurons=[10 25 50 100 150]; %matrix of how many neurons to use on each iteration
repetitions=50; %how many times to run the analysis
folds = 5; %number of times cross-validated:
shuffs = 1; %number of shuffled models created

%load parameters
BinDura=R_2R.Param.BinDura;
bins=R_2R.Param.bins;
binint=R_2R.Param.binint;
binstart=R_2R.Param.binstart;

%Need to keep number of trials used in VP and NAc constant, so have to pick
%minimum number across all sessions in both regions. Right now with
%existing sessions it's 20 Ev1s and 20 Ev2s
Ev1s=20;
Ev2s=20;



%find number of trials in each session
for i=1:length(RAW)
    
    %events being compared
    Ev1=strcmp('RD1', RAW(i).Einfo(:,2));
    Ev2=strcmp('RD2', RAW(i).Einfo(:,2));
    
    Ev1perSess(i)=length(RAW(i).Erast{Ev1,1});
    Ev2perSess(i)=length(RAW(i).Erast{Ev2,1});
end

%setup variables
PoolDec=[];
NN = 0;

%if you change this to e=1:3, you can also look at nonselective neurons
for e=1:2 %different selections of neurons

    
    %pick which set of neurons: all, reward-specific, or non-reward specific
    if e==1 selection=R_2R.Ninfo; end %all neurons
    if e==2 selection=R_2R.RSinfo; end %reward-selective neurons
    if e==3 selection=R_2R.RNSinfo; end %reward-nonselective neurons
    
    for f=1:2 %region

        TotalNeurons = 0;

        %total number of neurons in all sessions and events per session
        for i=1:length(RAW)
            if strcmp(region(f),RAW(i).Type(1:2))
                TotalNeurons=TotalNeurons+size(RAW(i).Nrast,1);
            end
        end

        for l=1:bins
            DecodeSpikes=NaN(1,1);
            AllSpikes=NaN(1,1);
            NN = 0;
            for i=1:length(RAW) %loops through sessions
                if strcmp(region(f),RAW(i).Type(1:2)) && Ev1perSess(i)>=20 && Ev2perSess(i)>=20
                    
                    %events being compared
                    Ev1=strcmp('RD1', RAW(i).Einfo(:,2));
                    Ev2=strcmp('RD2', RAW(i).Einfo(:,2));
                    
                    for j= 1:size(RAW(i).Nrast,1) %Number of neurons per session
                        if sum(strcmp(RAW(i).Ninfo(j,1),selection(:,1)) & strcmp(RAW(i).Ninfo(j,2),selection(:,2)))==1 %check if this neuron is on the selection list
                            NN=NN+1; %neuron counter

                            Ev1Spikes=NaN(1,1);
                            Ev2Spikes=NaN(1,1);

                            %get number of spikes for this neuron in this bin for all
                            %Ev1 trials
                            [PSR1,N1]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{Ev1},[(BinDura(1)+(binstart - binint)+l*binint) (BinDura(2)+(binstart - binint)+l*binint)],{2});% makes trial by trial rasters. PSR1 is a cell(neurons, trials)
                            for m=1:length(PSR1)
                                Ev1Spikes(m,1)=sum(PSR1{1,m}>(binstart));
                            end

                            %pick which trials get used for decoding
                            SetupTrials=cat(1,ones(Ev1s,1),zeros(length(PSR1)-Ev1s,1));
                            Trials=(SetupTrials(randperm(length(SetupTrials)))==1);

                            %put spikes from those trials in a matrix
                            AllSpikes(1:Ev1s,NN)=Ev1Spikes(Trials,1);

                            %get all the spikes from reward 1p2 trials
                            [PSR2,N2]=MakePSR04(RAW(i).Nrast(j),RAW(i).Erast{Ev2},[(BinDura(1)+(binstart - binint)+l*binint) (BinDura(2)+(binstart - binint)+l*binint)],{2});% makes trial by trial rasters. PSR1 is a cell(neurons, trials)
                            for n=1:length(PSR2)
                                Ev2Spikes(n,1)=sum(PSR2{1,n}>(binstart));
                            end

                            %pick which trials get used for decoding
                            SetupTrials2=cat(1,ones(Ev2s,1),zeros(length(PSR2)-Ev2s,1));
                            Trials2=(SetupTrials2(randperm(length(SetupTrials2)))==1);

                            %put spikes from those trials in a matrix
                            AllSpikes((Ev1s+1):(Ev1s+Ev2s),NN)=Ev2Spikes(Trials2,1);
                        end
                    end %neurons
                end %checking if enough events in that session
            end %sessions

            TotalNeurons=NN;    
            
            %do the decoding
            for v = 1:length(NumNeurons)
                if TotalNeurons>NumNeurons(v)
                    for u=1:repetitions
                        LowHz=zeros(1,NumNeurons(v)); %reset the lowHz identifier

                        %pick which neurons to use
                        SetupSel=cat(1,ones(NumNeurons(v),1),zeros(TotalNeurons-NumNeurons(v),1));
                        Sel=(SetupSel(randperm(length(SetupSel)))==1);

                        %setup the event identity matrix for decoding
                        %DecodeRs=zeros(Ev1s+Ev2s,NumNeurons(v));
                        DecodeRs(1:Ev1s,1)=1;
                        %DecodeRs(1:Ev1s,end)=1;
                        %DecodeRs((Ev1s+1):(Ev1s+Ev2s),2:end-1)=1;
                        DecodeRs((Ev1s+1):(Ev1s+Ev2s),1)=0;
                    
                        %setup decode spike matrix
                        DecodeSpikes=AllSpikes(:,Sel);
                        
                        %find neurons with too few spikes

                        for t=1:NumNeurons(v)
                            if sum(DecodeSpikes(1:Ev1s,t))<7 || sum(DecodeSpikes((1+Ev1s):(Ev1s+Ev2s),t))<7
                                LowHz(1,t)=1;
                            end
                        end

                        %remove neurons with too few spikes
                        if sum(LowHz)>0
                            Sel2=LowHz<1;
                            DecodeSpikes=DecodeSpikes(:,Sel2);
                        end

                        %set up validation result matrices
                        CVacc = NaN(folds,1);
                        CVaccSh = NaN(folds,1);

                        %normal model
                        Partitions = cvpartition(DecodeRs,'KFold',folds);
                        for r = 1:folds

                            SVMModel = fitcdiscr(DecodeSpikes(Partitions.training(r),:),DecodeRs(Partitions.training(r)));
                            prediction = predict(SVMModel,DecodeSpikes(Partitions.test(r),:));
                            actual = DecodeRs(Partitions.test(r));
                            correct = prediction - actual;
                            CVacc(r) = sum(correct==0) / length(correct);

                        end

                        %shuffled model
                        for q=1:shuffs
                            DecodeRsSh=DecodeRs(randperm(length(DecodeRs)));
                            PartitionsSh = cvpartition(DecodeRsSh,'KFold',folds);
                            for s = 1:folds
                                SVMModelSh = fitcdiscr(DecodeSpikes(PartitionsSh.training(s),:),DecodeRsSh(PartitionsSh.training(s)));
                                predictionSh = predict(SVMModelSh,DecodeSpikes(PartitionsSh.test(s),:));
                                actualSh = DecodeRsSh(PartitionsSh.test(s));
                                correctSh = predictionSh - actualSh;
                                CVaccSh(s) = sum(correctSh==0) / length(correctSh);
                            end
                            AccShuff(q,1) = nanmean(CVaccSh);
                        end

                        PoolDec{e,f}.True{v,1}(u,l) = nanmean(CVacc);
                        PoolDec{e,f}.Shuff{v,1}(u,l) = nanmean(AccShuff);

                        
                    end %repetitions
                    fprintf(['Selection #' num2str(e) ', Region #' num2str(f) ', Bin #' num2str(l) ', Condition #' num2str(v) '\n']);
                else
                    PoolDec{e,f}.True{v,1} = NaN;
                    PoolDec{e,f}.Shuff{v,1} = NaN;
                    fprintf(['Selection #' num2str(e) ', Region #' num2str(f) ', Bin #' num2str(l) ', Condition #' num2str(v) '\n']);
                end %checking if there are enough neurons
            end %conditions (number of neurons)
        end %bins
    end %region
end %selections

save('PoolDec_2R.mat','PoolDec');

R_2R.Param.NumNeurons=NumNeurons;
save('R_2R.mat','R_2R');