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

%%%% Analyse Network Choices


function analyse_TS_Learning_InitCluster


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;
disp([num2str(Ns),' batches.'])


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

%%%Initial learning: C0  0 1  / 1 2 % 01
%%%Initial learning: C1  0 1  / 1 2 % 23
%%%Initial learning: C2  2 3  / 3 4 % 45



%%%%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
%%%%%%%%%%%%%%%%%%%%%

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

NTransfer=length(find(Stim==6|Stim==7))/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=0;%length(find(Stim<4))/4-NLearn;%%%Nb of stimulus occurences in noisy asymptotic learning phase

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

%%%Set reward values
SSE=X(:,5);
SSE(SSE<.5)=0;
SSE(SSE>=.5)=1;
%%%

Su=1:Ns;


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

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

mi = min(NLearn);
ma = max(NLearn);

Su(NLearn==100)=[];
disp([num2str(length(find(NLearn >=99))), ' networks out of ',...
    num2str(Ns),' did not learn within 100 epochs. Taken out of further analysis.'])
mi = min(NLearn(Su));
ma = max(NLearn(Su));

disp(['Criterion reached within [',num2str(mi),',',num2str(ma),'] epochs.'])

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)-5),'facecolor',[1 1 1]);
hold on
errorbar(1, mean(NLearn(Su)-5),std(NLearn(Su))/sqrt(N),'k');
hold off
title('Initial Learning Criterion')
ylabel('Epochs')
%boum


stim=[0 2 4];
%%%%Laerning phases learning curves
Ns=length(Su);
DonneesM=[];
DonneesS=[];
Asym=[];
N7 =7 ;
for i=[1:3]%%%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));  
        bla = bla';%mean(bla(1:10))
        resume(N,i) = 1-mean(bla(1:N7));%
        if length(bla)<152
            bla=[bla bla(end)*ones(1,152-length(bla))];
        end
        Y=[Y;bla];
        Donnees(N,:,i)=bla;
    end
    DonneesM=[DonneesM;mean(Y)];
    DonneesS=[DonneesS;std(Y)/sqrt(size(Y,1))];% m=mean(Final);
   
end
for N=Su
 figure(3)
    hold on
    plot(1-resume(N,:),'color',rand(1,3))
end
errorbar(1-mean(resume(Su,[1 2 3])),std(resume(Su,[1 2 3]))/sqrt(length(N)),'k')
title('First 7 trials error proportion')
set(gca,'xtick',1:3,'xticklabel',{'C0TS1','C1TS1','C2TS2'})


%%% Resume: proportion of errors on first N7 trials per (context x stim)
[h,p,ci,stats]=ttest(mean(resume(Su,[1 2]),2)-resume(Su,3));
disp(['% Whole group: transfer: N = ',num2str(length(Su)),' ; t = ',num2str(stats.tstat),' ; p = ',num2str(p)])

types={'b','g','r'};
legendes={'C0 TS1','C1 TS1','C2 TS2'}; 

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


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


P1=[];
P1s=[];
P2=[];
P2s=[];
P3=[];
P3s=[];
P4=[];
P4s=[];
PP1=[];
PP2=[];
PP3=[];
for N=Su 
    Prop1=[];
    Prop2=[];
    Prop3=[];
    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
    TS3=X(X(:,1)==N & [X(:,4)==4 |X(:,4)==5],7);%%%Sequence of chosen TS for Context 3
    
    TS1f=TS1(2*NLearn(N)-10+1:2*NLearn(N));%%%Restricted to last ten trials of learning phase (ie when no error)
    TS2f=TS2(2*NLearn(N)-10+1:2*NLearn(N));
    TS3f=TS3(4*NLearn(N)-10+1:4*NLearn(N));
    L1=[length(find(TS1f==1)) length(find(TS1f==2)) length(find(TS1f==3))];
    L2=[length(find(TS2f==1)) length(find(TS2f==2)) length(find(TS2f==3))];
    L3=[length(find(TS3f==1)) length(find(TS3f==2)) length(find(TS3f==3))];
    ts1=find(L1==max(L1));%%stripe number for C1 TS
    ts2=find(L2==max(L2));%%stripe number for C2 TS2
    ts3=find(L3==max(L3));%%stripe number for C3 TS3
    %[ts1 ts2 ts3]
    
    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))];
    L3=[length(find(TS3==1)) length(find(TS3==2)) length(find(TS3==3))];

    if ts1==ts2
        proportion(N)=(L1(ts1)+L2(ts2))/(4*NLearn(N));
    else
        proportion(N)=-(min(L1(ts1),L1(ts2))+min(L2(ts1),L2(ts2)))/(4*NLearn(N));
    end


end


% get non clustering (ts(C0) ~= ts(C1)) networks
si=find(proportion<0);si1=si;
Zi=-mean(resume(si,[1:2]),2)+resume(si,[3]);

[h p ci stats]=ttest(Zi);[p stats.tstat length(si)];
disp(['% No cluster group: transfer: N = ',num2str(length(si)),' ; t = ',num2str(stats.tstat),' ; p = ',num2str(p)])

figure(figu-2)
subplot(1,2,1)
for i=[1 2 3]%1:8
    plot([1:100],1-mean(Donnees(si,1:100,i)),types{i})
    hold on
end
legend(legendes)
title('Training phase learning curves no cluster')
for i=[1 2 3]%1:8
    tracerErreur([1:100],1-mean(Donnees(si,1:100,i)),std(Donnees(si,1:100,i))/sqrt(length(si)),types{i})
    hold on
end
hold off



% get clustering (ts(C0) == ts(C1)) networks
si=find(proportion>0);si2=si;
Zi=-mean(resume(si,[1:2]),2)+resume(si,[3]);

[h p ci stats]=ttest(Zi);[p stats.tstat length(si)];
disp(['% Cluster group: transfer: N = ',num2str(length(si)),' ; t = ',num2str(stats.tstat),' ; p = ',num2str(p)])


    figure(figu-2)
    subplot(1,2,2)
for i=[1 2 3]%1:8
    plot([1:100],1-mean(Donnees(si,1:100,i)),types{i})
    hold on
end
legend(legendes)
title('Training phase learning curves cluster')
for i=[1 2 3]%1:8
    tracerErreur([1:100],1-mean(Donnees(si,1:100,i)),std(Donnees(si,1:100,i))/sqrt(length(si)),types{i})
    hold on
end
hold off

Z=+mean(resume(si,[1:2]),2)-resume(si,[3]);
figure(97)
plot(proportion(si),Z','.')


% relate increase in perf to clustering in TS space
[rho pval]=corr((proportion(si))',+mean(resume(si,[1:2]),2)-resume(si,[3]));
disp(['% Cluster group: TS vs transfer: rho = ',num2str(rho),' ; p = ',num2str(pval)])
a=rho*std(Z)/std(proportion(si));
b=mean(Z)-a*mean(proportion(si));
hold on
plot([.35:.01:1],a*[.35:.01:1]+b,'k')
hold off
xlabel('early TS clustering choice')
ylabel('early TS1 vs TS2 performance benefit')
title('Performance benefit depends on initial clustering')

[rho pval]=corr((proportion(si))',resume(si,[3]));
disp(['% Cluster group: TS vs TS2: rho = ',num2str(rho),' ; p = ',num2str(pval)])
[rho pval]=corr((proportion(si))',mean(resume(si,[1:2]),2));
disp(['% Cluster group: TS vs TS1: rho = ',num2str(rho),' ; p = ',num2str(pval)])

Z1=-mean(resume(si1,[1:2]),2)+resume(si1,[3]);
Z2=-mean(resume(si2,[1:2]),2)+resume(si2,[3]);
[h p ci stats]=ttest2(Z1,Z2);
disp(['% Comparaison gr transfer:  t = ',num2str(stats.tstat),' ; p = ',num2str(p)])

Z1=resume(si1,[3]);
Z2=resume(si2,[3]);
[h p ci stats]=ttest2(Z1,Z2);
disp(['% Comparaison gr TS2:  t = ',num2str(stats.tstat),' ; p = ',num2str(p)])
Z1=mean(resume(si1,[1:2]),2);
Z2=mean(resume(si2,[1:2]),2);
[h p ci stats]=ttest2(Z1,Z2);
disp(['% Comparaison gr TS1:  t = ',num2str(stats.tstat),' ; p = ',num2str(p)])



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

