% *** MATLAB CODE. CHANGE EXTENSION TO .m BEFORE EXECUTING ***
%
% This program calculates the probability of each possible outcome when a
% group is given two options.
% SYNTAX:
% [prob_final,n_y,dynamics]=ProtocolS1(N_tot,a0,a1,a2,a3,a1_replicas,a2_replicas,replicas_xy)
% All arguments are optional. You can try it with no parameters, and obtain
% the results for the case of 8 individuals in the symmetric set-up with
% a1=2.5.
% - N_tot is the number of inviduals in the group. Default: N_tot=8
% - a0 is the parameter corresponding to non-social information. Default: a0=1
% - a1 is the parameter measuring the reliability of other individuals. It
% affects to all individuals if a2 is not specified. If a2 is specified, it
% affects only to individuals choosing x. Default: a1=2.6
% - a2 is the parameter measuring the reliability of individuals choosing y. Default: a2=a1.
% - a3 is the parameter associated to the information coming from undecided
% individuals. Default: a3=1.
% - a1_replicas is the same as a1, but only applies to the replicas. Default: a1_replicas=a1.
% - a2_replicas is the same as a2, but only applies to the replicas. Default: a2_replicas=a2.
% - replicas_xy is a two-element vector with the number of conspecific
% replicas choosing [x y]. Default: replicas_xy=[0 0].
%
% - prob_final is the probability of each final state. It is a vector with
% N_tot+1 elements. Each element corresponds to the state with the number
% of individuals going to y that is specified in vector n_y. 
% - n_y is a vector with N_tot+1 elements. It contains the number of
% individuals going to y that correspond to each state in prob_final.
% - dynamics is a N_tot+1 x N_tot+1 matrix, that contains the probability
% of each state as each individual makes its decision. The
% diagonal from the bottom-left to the upper-right corners of this matrix
% contains the same information as prob_final.
%
% Alfonso Pérez-Escudero & Gonzalo G. de Polavieja. Instituto Cajal,
% Consejo Superior de Investigaciones Científicas, Madrid, Spain
% alfonso.perez.escudero@gmail.com, gonzalo.polavieja@cajal.csic.es
% Last update: July 30, 2010
%
% Cite as: Alfonso Pérez-Escudero, Gonzalo G. de Polavieja (2011)
% Collective Animal Behavior from Bayesian Estimation and Probability
% Matching, PLoS Computational Biology

function [prob_final,n_y,dynamics]=ProtocolS1(N_tot,a0,a1,a2,a3,a1_replicas,a2_replicas,replicas_xy)

% Default values for the parameters
if nargin<1 || isempty(N_tot)
    N_tot=8;
end
if nargin<2 || isempty(a0)
    a0=1;
end
if nargin<3 || isempty(a1)
    a1=2.5;
end
if nargin<4 || isempty(a2)
    a2=a1;
end
if nargin<5 || isempty(a3)
    a3=1;
end
if nargin<6 || isempty(a1_replicas)
    a1_replicas=a1;
end
if nargin<7 || isempty(a2_replicas)
    a2_replicas=a2;
end
if nargin<8 || isempty(replicas_xy)
    replicas_xy=[0 0];
end


nbichos_L=repmat((0:N_tot)',[1 N_tot+1]);
nbichos_R=nbichos_L';

prob_R=1./(1 + a0*a1.^nbichos_L.*a1_replicas^replicas_xy(1).*a2.^-nbichos_R*a2_replicas^-replicas_xy(2).*a3.^(N_tot-nbichos_L-nbichos_R-1));

dynamics=zeros(N_tot+1,N_tot+1);

dynamics(1,1)=1;
for c_bichos=1:N_tot
    estados_act = c_bichos:N_tot:c_bichos + N_tot*(c_bichos-1);
    siguientes_R = estados_act+N_tot+1; % Goes to the next element to the right
    siguientes_L = estados_act+1; % Goes to the next element downwards
    dynamics(siguientes_R)=dynamics(estados_act).*prob_R(estados_act);
    dynamics(siguientes_L) = dynamics(siguientes_L) + dynamics(estados_act).*(1-prob_R(estados_act));  
end % c_bichos
n_y=0:N_tot;
c_bichos=c_bichos+1;
estados_act = c_bichos:N_tot:c_bichos + N_tot*(c_bichos-1);
prob_final=dynamics(estados_act); 

% *********************************
% * REPRESENTATION OF THE RESULTS *
% *********************************

figure
plot(n_y,prob_final)
xlabel('Number of individuals going to \ity')
ylabel('Probability')
title('Final outcome')
figure
imagesc(n_y,n_y,dynamics)
xlabel('Number of individuals going to \ity')
ylabel('Number of individuals going to \itx')
title('Dynamics')