%%%% Anne Collins
%%%% LNCC, Brown University
%%%% april 2011
%%%% project: Emergent TS_Learning.proj

%%%% Analyse Network Choices


function analyse_TS_Learning


figu=100;%floor(100*rand)+1;

%%%Load network trial by trial data
load TS_Learning_Example.mat %%% Data used in paper Collins & Frank submitted
%X=csvread('TS_Learn.csv'); %%% Use that to analyse your own output from TS_Learning.proj

%%%Get number of batches
Ns=max(X(:,1))+1


%%%%Experimental phase/ Context / Stimuli numbers / correct actions numbers

%%%Initial learning: C1  0 1  / 1 2
%%%Initial learning: C2  2 3  / 3 4
%%%Transfert :       C3  4 5  / 1 2
%%%New Overlap :     C4  6 7  / 1 4
%%%New Incongruent : C5  8 9  / 2 3



%%%%Columns of X
%%% 1: batch
%%% 2: epoch
%%% 3: trial
%%% 4: input number
%%% 5: sse
%%% 6: RT
%%% 7: selected TS
%%% 8: selected action




%%%%%%%%%%%%%%%%%%%%%
%%% Determine different phases lengths
%%%%%%%%%%%%%%%%%%%%%


Stim=X(X(:,1)==0,4);%%%Stimulus sequence for first batch

NTransfer=length(find(Stim==4|Stim==5))/2;%%%Number of stimulus occurences in transfer test phase
T1=find(Stim>3);T1=T1(1);%%%1st time after learning period
NLearn=(T1-1)/4;%%%Nb of stimulus occurences in learning phase
NNoise=length(find(Stim<4))/4-NLearn;%%%Nb of stimulus occurences in noisy asymptotic learning phase
%boum



%%%%%%%%%%%%%%%%%%%%%
%%% Analyses
%%%%%%%%%%%%%%%%%%%%%

%%%Set reward values
SSE=X(:,5);
SSE(SSE<.5)=0;
SSE(SSE>.4)=1;
%%%
stim=[0 2 4 14 6 16 8 18];

Su=0:Ns-1;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%              Performances              %%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%Learning phase length
for N=[0:Ns-1]
    Stim=X(X(:,1)==N,4);
    NLearn(N+1)=length(find(Stim<4))/4-1*NNoise;%%%Nb de fois par stimuli ds learning phase
end

Su=[0:Ns-1];
if find(NLearn<5)
    disp(['Problem in Learning phase length'])
    boum
end
types={'k','k--'};
legendes={'TS1 = 1 2','TS2 = 3 4','Transfert TS1','NewO 1 4','NewI 2 3'};% 


figure(figu)
bar(1, mean(NLearn(Su+1)-5),'facecolor',[1 1 1]);
hold on
errorbar(1, mean(NLearn(Su+1)-5),std(NLearn(Su+1))/sqrt(N),'k');
hold off
title('Initial Learning Criterion')
%boum

%%%%Test phases learning curves
Ns=length(Su);
DonneesM=[];
DonneesS=[];
Asym=[];
for i=[3 5 7]%%%3 test phases
    Y=[];
    for N=Su %%%Batches considered
        %%%Performance for this batch over both stimuli of the considered
        %%%phase
        bla=.5 * (SSE(X(:,1)==N & X(:,4)==stim(i)) + SSE(X(:,1)==N & X(:,4)==stim(i)+1));        
        Y=[Y;bla'];
    end
    DonneesM=[DonneesM;mean(Y)];
    DonneesS=[DonneesS;std(Y)/sqrt(size(Y,1))];% m=mean(Final);
end

types={'b','g','r'};
legendes={'Transfert TS1','New Overlap 1 4','New Inc 2 3'}; 

for i=[1 2 3]%1:8
    figure(figu+1)
    plot([1:NTransfer],1-DonneesM(i,1:NTransfer),types{i})
    hold on
end
legend(legendes)
title('Test phase learning curves')
for i=[1 2 3]%1:8
    figure(figu+1)
    tracerErreur([1:NTransfer],1-DonneesM(i,1:NTransfer),DonneesS(i,1:NTransfer),types{i})
    hold on
end
hold off

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5
%%%%%%% Chosen Task-Sets
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


P1=[];
P1s=[];
P2=[];
P2s=[];
P3=[];
P3s=[];
P4=[];
P4s=[];
PP1=[];
PP2=[];
PP3=[];
for i=[3 5 7]%1:81 2 
    Prop1=[];
    Prop2=[];
    Prop3=[];
    for N=Su
        TS1=X(X(:,1)==N & [X(:,4)==0 |X(:,4)==1],7);%%%Sequence of chosen TS for Context 1
        TS2=X(X(:,1)==N & [X(:,4)==2 |X(:,4)==3],7);%%%Sequence of chosen TS for Context 2

        TS1=TS1(2*NLearn(N+1)-10+1:2*NLearn(N+1));%%%Restricted to last ten trials of learning phase (ie when no error)
        TS2=TS2(2*NLearn(N+1)-10+1:2*NLearn(N+1));
        L1=[length(find(TS1==1)) length(find(TS1==2)) length(find(TS1==3))];
        L2=[length(find(TS2==1)) length(find(TS2==2)) length(find(TS2==3))];
        L1=find(L1==max(L1));%%stripe number for C1 TS
        L2=find(L2==max(L2));%%stripe number for C2 TS2
        %%% ensure the task-set is well defined
        if length(L1)>1 & length(L2)>1
            L1=L1(1);
            L2=L1(2);
        elseif length(L1)>1
            L1(L1==L2)=[];
            L1=L1(1);
        elseif length(L2)>1
            L2(L2==L1)=[];
            L2=L2(1);
        end
        %%% rename task-sets so that TS1 is C1, TS2 corresponds to C2, TS3
        %%% to the non chosen one
        TS=X(X(:,1)==N,7);
        TS(TS==L1)=11;%%%TS utilise pour C1
        TS(TS==L2)=12;%%%TS utilise poour C2
        TS(TS<10)=13;%%% autre TS

        TS=TS-10;

        %%% Get chosen task-set sequence for current test phase
        I=X(X(:,1)==N,4); 
        bla=TS(I==stim(i) |I==stim(i)+1);
        bla=bla';
        
        %%%Binary vectors indicating when each TS is chosen
        proportion1=0*bla;
        proportion1(bla==1)=1;
        proportion2=0*bla;
        proportion2(bla==2)=1;
        proportion3=0*bla;
        proportion3(bla==3)=1;
        
        %%%accumulate
        Prop1=[Prop1;proportion1];
        Prop2=[Prop2;proportion2];
        Prop3=[Prop3;proportion3];
        
    end
    
    %%%Get initial average use of each task-set
    PP1=[PP1 mean(Prop1(:,1:floor(end/5)),2)];
    PP2=[PP2 mean(Prop2(:,1:floor(end/5)),2)];
    PP3=[PP3 mean(Prop3(:,1:floor(end/5)),2)];

    %%%Get learning curves of use of each TS, mean and sem
    P1=[P1;mean(Prop1)];
    P1s=[P1s;std(Prop1)/sqrt(Ns)];
    P2=[P2;mean(Prop2)];
    P2s=[P2s;std(Prop2)/sqrt(Ns)];
    P3=[P3;mean(Prop3)];
    P3s=[P3s;std(Prop3)/sqrt(Ns)];
    P4=[P4;mean(mean(Prop3-(Prop1+Prop2)))];
    P4s=[P4s;std(mean(Prop3-(Prop1+Prop2),2))/sqrt(Ns)];
    if i==3
        PosTr=mean(Prop1,2);
        NegTr=mean(Prop3,2);
    elseif i==5
        %NegTr=mean(Prop3,2);
    end
    
end

%%%Plot early average TS use as a function of condition
figure(12)
errorbar(mean(PP1),std(PP1)/sqrt(size(PP1,1)),'c')
hold on
errorbar(mean(PP2),std(PP2)/sqrt(size(PP2,1)),'m')
errorbar(mean(PP3),std(PP3)/sqrt(size(PP3,1)),'y')
hold off
Title('Initial TS use')
Legend({'TS1','TS2','TS3'})
set(gca,'xtick',[1:3])
set(gca,'xticklabel',{'Transfer','NewO','NewInc'})

%%%
base=[1:NTransfer];
figure(figu+4)
for i=1:3
    subplot(1,3,i)
    tracerErreur(base,P1(i,1:NTransfer),P1s(i,1:NTransfer),'c')
    hold on
    tracerErreur(base,P2(i,1:NTransfer),P2s(i,1:NTransfer),'m')
    tracerErreur(base,P3(i,1:NTransfer),P3s(i,1:NTransfer),'y')    
    title(legendes{i})
end


%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% Analysis of test error type
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


TS=[1 2 3 4 1 2 1 4 2 3];%%% Correct actions
GC=[3 4 1 2 3 4 3 2 1 4];%%% Neglect High dim errors
GS=[2 1 4 3 2 1 4 1 3 2];%%% Neglect Low dim errors
GE=[4 3 2 1 4 3 2 3 4 1];%%% Neglect all errors

Erreurs=[];

for N=Su
    %%%Get batch stimuli and actions for test phases
    Actions=X(X(:,1)==N,8);
    Stimuli=1+X(X(:,1)==N,4);
    Actions=Actions(Stimuli>4);
    Stimuli=Stimuli(Stimuli>4);

    %%%Build sequences of errortypes
    GenC=0*Actions;
    for s=5:10%[1 3]%
        GenC(Actions==GC(s) &Stimuli==s)=1;
    end
    GenS=0*Actions;
    for s=5:10%[1 3]%
        GenS(Actions==GS(s) &Stimuli==s)=1;
    end
    Expl=0*Actions;
    for s=5:10%[1 3]%
        Expl(Actions==GE(s) &Stimuli==s)=1;
    end
    Cor=0*Actions;
    for s=5:10
        Cor(Actions==TS(s)&Stimuli==s)=1;
    end

    %%%Get error repartitions
    bla=[];
    for jj=[5 7 9]
        GeC=GenC(Stimuli==jj|Stimuli==jj+1);
        GeS=GenS(Stimuli==jj|Stimuli==jj+1);
        GeE=Expl(Stimuli==jj|Stimuli==jj+1);
        Co=Cor(Stimuli==jj|Stimuli==jj+1);
        %[GC(1:10)';GS(1:10)';GE(1:10)';Co(1:10)']

        bla=[bla mean(GeC(1:10)) mean(GeS(1:10)) mean(GeE(1:10))];
    end
    Erreurs=[Erreurs;bla];
end


figure(figu+6)
errorbar(mean(Erreurs(:,1:3)),std(Erreurs(:,1:3))/sqrt(Ns),'b')
hold on
errorbar(mean(Erreurs(:,4:6)),std(Erreurs(:,4:6))/sqrt(Ns),'g')%[.85:1:2.85],
errorbar(mean(Erreurs(:,7:9)),std(Erreurs(:,7:9))/sqrt(Ns),'r')%1.15:1:3.15],
hold off
Title('Test Phase Error Repartition')
Legend(legendes)
set(gca,'xtick',[1:3])
set(gca,'xticklabel',{'NH','NL','NA'})

%boum

%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% Analysis of asymptotic phase errors
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%Errortypes
TS=[1 2 3 4];
GC=[3 4 1 2];
GS=[2 1 4 3];
GE=[4 3 2 1];


%%%Get them
Erreurs=[];
RTEr=[];
for N=Su
    %%%Get actions and input sequences sequences
    Actions=X(X(:,1)==N,8);
    Stimuli=1+X(X(:,1)==N,4);
    %%%Get context, previous context
    Cont=floor(X(X(:,1)==N,4)/2);
    CTp1=[Cont(1) ;Cont(1:end-1)];
    %%%Determine if this is a switch or stay trial
    Sw=0*Cont;
    Sw(CTp1==Cont)=1;
    Sw=1-Sw;
    ta=length(Actions);
    %%%Get chosen Task-set and reaction times
    TS=X(X(:,1)==N ,7);
    RT=X(X(:,1)==N ,6);
    %%%Restrict all to noisy asymptotic phase
    Actions=Actions(end-4*NNoise+1:end);
    Stimuli=Stimuli(end-4*NNoise+1:end);
    TS=TS(end-4*NNoise+1:end);
    Sw=Sw(end-4*NNoise+1:end);
    RT=RT(end-4*NNoise+1:end);
    %%%Get error types and their reaction times
    GenC=0*Actions;
    RTC=[];
    for s=1:4%[1 3]%
        GenC(Actions==GC(s) &Stimuli==s & TS<3 & Sw==1)=1;
        GenC(Actions==GC(s) &Stimuli==s & TS<3 & Sw==0)=-1;
        RTC=[RTC ;RT(Actions==GC(s) &Stimuli==s & TS<3& Sw==1)];%
    end
    GenS=0*Actions;
    RTS=[];
    for s=1:4%[1 3]%
        GenS(Actions==GS(s) &Stimuli==s & TS<3 & Sw==1)=1;
        GenS(Actions==GS(s) &Stimuli==s & TS<3 & Sw==0)=-1;
        RTS=[RTS ;RT(Actions==GS(s) &Stimuli==s & TS<3& Sw==1)];%
    end
    Expl=0*Actions;
    RTE=[];
    for s=1:4%[1 3]%
        Expl(Actions==GE(s) &Stimuli==s & TS<3 & Sw==1)=1;
        Expl(Actions==GE(s) &Stimuli==s & TS<3 & Sw==0)=-1;
        RTE=[RTE ;RT(Actions==GE(s) &Stimuli==s & TS<3 & Sw==1)];%
    end
    %%%Number of Switch, Stay trials
    nErSw=length(Sw==1 & TS<3);%
    nErSt=length(Sw==0 & TS<3);%
    nEr=1;
    %%% Get proportion of different kinds of erros, depending on switch or
    %%% stay
    if nErSw==0
        bla=[0 0 0];
    else
        bla=[length(find(GenC==1 & TS<3))*nEr/nErSw];%length(GenC)];%
        bla=[bla length(find(GenS==1 & TS<3))*nEr/nErSw];%length(GenS)];%
        bla=[bla length(find(Expl==1 & TS<3))*nEr/nErSw];%length(Expl)];%
    end
    if nErSt==0
        bla=[bla 0 0 0];
    else
        bla=[bla length(find(GenC==-1 & TS<3))*nEr/nErSt];%length(GenC)];%
        bla=[bla length(find(GenS==-1 & TS<3))*nEr/nErSt];%length(GenS)];%
        bla=[bla length(find(Expl==-1 & TS<3))*nEr/nErSt];%length(Expl)];%
    end
    Erreurs=[Erreurs;bla];
    
    %%% reaction times of different kinds of errors
    if isempty(RTC)
        bla=[-1];%%% signal if there are no errors of that type
    else
        bla=mean(RTC);
    end
    if isempty(RTS)
        bla=[bla -1];%%% signal if there are no errors of that type
    else
        bla=[bla mean(RTS)];
    end
    if isempty(RTE)
        bla=[bla -1];%%% signal if there are no errors of that type
    else
        bla=[bla mean(RTE)];
    end
    RTEr=[RTEr;bla];
end
Xe=Erreurs(:,2)-Erreurs(:,1);
m=mean(Xe);
s=std(Xe);
 
 
RTErC=RTEr(:,1); 
RTErS=RTEr(:,2); 
RTErE=RTEr(:,3); 
% RTErC(RTErE==-1)=[];
% RTErS(RTErE==-1)=[];
RTErE(RTErE==-1)=[];
% RTErE(RTErS==-1)=[];
% RTErC(RTErS==-1)=[];
RTErS(RTErS==-1)=[];
% RTErS(RTErC==-1)=[];
% RTErE(RTErC==-1)=[];
RTErC(RTErC==-1)=[];
NNs=[length(RTErC) length(RTErS) length(RTErE)];
%[h p ci stat]=ttest2(RTErS,RTErC);
% 
%  %Erreurs
figure(figu+7)
hold on
errorbar(mean(Erreurs(Su+1,1:3)),std(Erreurs(Su+1,1:3))/sqrt(Ns),'k')%[.85:1:2.85],
errorbar(mean(Erreurs(Su+1,4:6)),std(Erreurs(Su+1,4:6))/sqrt(Ns),'color',[.5 .5 .5])%1.15:1:3.15],
hold off
Title('Noisy asymptotic Learning phase Errors')
Legend({'Switch','Stay'})
set(gca,'xtick',[1:3])
set(gca,'xticklabel',{'NH','NL','NA'})
figure(figu+16)
bar([mean(RTErC) mean(RTErS)],'facecolor',[1 1 1])
hold on
errorbar([mean(RTErC) mean(RTErS)],[std(RTErC) std(RTErS)]./sqrt(NNs(1:2)),'+k')
hold off
Title('Noisy asymptotic Learning phase Errors Reaction Times')
set(gca,'xtick',[1:2])
set(gca,'xticklabel',{'NH','NL'})

%boum

%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% Analyse switch cost
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



Co=[1 2 3 4];
Erreurs=[];
RTSw=[];
ERSw=[];
for N=Su
    %%%Get Actions
    Actions=X(X(:,1)==N,8);
    Stimuli=1+X(X(:,1)==N,4);
    %%%Get Context, Switch, Task set, RTs
    Cont=floor(X(X(:,1)==N,4)/2);
    CTp1=[Cont(1) ;Cont(1:end-1)];
    Sw=0*Cont;
    Sw(CTp1==Cont)=1;
    Sw=1-Sw;
    ta=length(Actions);
    TS=X(X(:,1)==N ,7);
    RT=X(X(:,1)==N ,6);
    %%%Limit to asymptotic learning phase
    Actions=Actions(end-4*NNoise+1:end);
    Stimuli=Stimuli(end-4*NNoise+1:end);
    TS=TS(end-4*NNoise+1:end);
    Sw=Sw(end-4*NNoise+1:end);
    RT=RT(end-4*NNoise+1:end);

    if find(Sw==1)
        RTw=[RT(Sw==1)];%
        RTt=[RT(Sw==0)];%

        Erw=0;
        Ert=0;
        for s=1:4%[1 3]%
            Erw=Erw+1-length(find(Sw==1 & Stimuli==s & TS<3 & Actions==Co(s) ))/length(find(Sw==1& TS<3 &Stimuli==s));%
            Ert=Ert+1-length(find(Sw==0 & Stimuli==s& TS<3 & Actions==Co(s)))/length(find(Sw==0& TS<3& Stimuli==s));%
        end
        Erw=Erw/4;
        Ert=Ert/4;
        RTSw=[RTSw [mean(RTw); mean(RTt); mean(RTw)-mean(RTt)]];
        ERSw=[ERSw Erw-Ert];%
    end
end
    
    
    
figure(figu+9)
    subplot(1,2,1)
bar(mean(RTSw(1:2,:)'),'facecolor',[1 1 1])
hold on
errorbar(mean(RTSw(1:2,:)'),std(RTSw(1:2,:)')/sqrt(Ns),'+k')
hold off
Title('Reaction Time Switch Cost')
set(gca,'xtick',[1:2])
set(gca,'xticklabel',{'Switch','Stay'})

figure(figu+9)
    subplot(1,2,2)
bar(mean(ERSw),'facecolor',[1 1 1])
hold on
errorbar(mean(ERSw),std(ERSw)/sqrt(Ns),'+k')
hold off
Title('Error Switch Cost')
set(gca,'xtick',[1])
set(gca,'xticklabel',{'Switch-Stay'})

end



function tracerErreur(Y,X1,X2,coul)


plot(Y,X1,coul,'linewidth',2)
hold on
plot(Y,X1+X2,coul,'linewidth',.5)
plot(Y,X1-X2,coul,'linewidth',.5)
%hold off
end

