% *** MATLAB CODE. CHANGE EXTENSION TO .m BEFORE EXECUTING ***
%
% This program calculates the probability of each possible outcome when a
% group is given two options, according to the model that takes into
% account dependencies among the other fish.
% SYNTAX:
% [prob_final,n_y,dynamics,dynamics_cell,states,dynamics_givenL,dynamics_givenR]=ProtocolS2(N_tot,a0,a1,a2,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=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.
% - 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. In this matrix all 
% states with the same n_x:n_y are collapsed into a single state. The
% diagonal from the bottom-left to the upper-right corners of this matrix
% contains the same information as prob_final. CAUTION: This matrix is only given
% for comparison with the results of groupsimulator.m. In this program, the
% relevant probabilities are in dynamics_cell.
% - dynamics_cell is a cell that contains the probability of each state.
% Each element of the cell refers to one step in the simulation, and
% contains a vector with the probabilities of all the states that correspond to
% that step.
% - states is a cell with the same dimensions as dynamics_cell. Each
% element contains an Mx2 matrix, where M is the number of states that 
% correspond to the current step. This matrix contains the labels n_x:n_y 
% of the states in dynamics_cell.
% - dynamics_givenL is a cell with the same dimensions as dynamics_cell,
% that gives the probability of each state, under the assumption that x is
% the best choice. This is only an auxiliary variable, used for the
% recursive calculations of S and the probabilities in dynamics_cell.
% - dynamics_givenR is the same as dynamics_givenR, but under the
% assumption that y is the best choice.
%
% 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,dynamics_cell,states,dynamics_givenL,dynamics_givenR]=ProtocolS2(N_tot,a0,a1,a2,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=5;
end
if nargin<4 || isempty(a2)
    a2=a1;
end
if nargin<5 || isempty(a1_replicas)
    a1_replicas=a1;
end
if nargin<6 || isempty(a2_replicas)
    a2_replicas=a2;
end
if nargin<7 || isempty(replicas_xy)
    replicas_xy=[0 0];
end



dynamics=cell(1,N_tot+1);
dynamics_givenL=cell(1,N_tot+1);
dynamics_givenR=cell(1,N_tot+1);
states=cell(1,N_tot+1);

dynamics{1}=1;
dynamics_givenL{1}=1;
dynamics_givenR{1}=1;
states{1}=[0 0];
% The model always starts from 0:0 The first iteration of this loop is to
% simulate the possible behaviors of the replicas (provided X or Y), that
% will be used to compute the initial condition for the second iteration,
% in which the real fish are simulated
for c_vueltas=1:2 
    if c_vueltas==1
        a1_act=a1_replicas;
        a2_act=a2_replicas;
        n_pasos=sum(replicas_xy);
    else
        a1_act=a1;
        a2_act=a2;
        n_pasos=N_tot;
    end
    for c_bichos=1:n_pasos
        states{c_bichos+1}(1:2:2^c_bichos,:)=states{c_bichos}+repmat([1 0],[2^(c_bichos-1) 1]);
        states{c_bichos+1}(2:2:2^c_bichos,:)=states{c_bichos}+repmat([0 1],[2^(c_bichos-1) 1]);
        probR_act=1./(1+a1_act*dynamics_givenL{c_bichos}./dynamics_givenR{c_bichos});
        dynamics_givenL{c_bichos+1}(2:2:2^c_bichos) = dynamics_givenL{c_bichos}.*probR_act;
        dynamics_givenL{c_bichos+1}(1:2:2^c_bichos) = dynamics_givenL{c_bichos}.*(1-probR_act);
        probR_act=1./(1+a2_act^-1*dynamics_givenL{c_bichos}./dynamics_givenR{c_bichos});
        dynamics_givenR{c_bichos+1}(2:2:2^c_bichos) = dynamics_givenR{c_bichos}.*probR_act;
        dynamics_givenR{c_bichos+1}(1:2:2^c_bichos) = dynamics_givenR{c_bichos}.*(1-probR_act);
        if c_vueltas==2
            probR_act=1./(1+a0*dynamics_givenL{c_bichos}./dynamics_givenR{c_bichos});
            dynamics{c_bichos+1}(2:2:2^c_bichos) = dynamics{c_bichos}.*probR_act;
            dynamics{c_bichos+1}(1:2:2^c_bichos) = dynamics{c_bichos}.*(1-probR_act);
        end    
    end % c_bichos
    if c_vueltas==1
        nivel=sum(replicas_xy)+1;
        buenos=states{nivel}(:,1)==replicas_xy(1) & states{nivel}(:,2)==replicas_xy(2);
        prob_givenL=sum(dynamics_givenL{nivel}(buenos));
        prob_givenR=sum(dynamics_givenR{nivel}(buenos));
        dynamics_givenL=cell(1,N_tot+1);
        dynamics_givenR=cell(1,N_tot+1);
        dynamics_givenL{1}=prob_givenL;
        dynamics_givenR{1}=prob_givenR;
    end
end % c_vueltas
n_y=0:N_tot;

% Transform to the usual dynamics matrix. This is only done for comparison
% with the results provided by gropusimulator.m, because for this case all
% states are kept separately in dynamics_cell
dynamics_cell=dynamics;
dynamics=zeros(N_tot+1);
for c_escalones=1:N_tot+1
    for c_casos=1:c_escalones
        estado_act=[c_casos-1 c_escalones-c_casos];
        buenos=states{c_escalones}(:,1)==estado_act(1) & states{c_escalones}(:,2)==estado_act(2);
        dynamics(estado_act(1)+1,estado_act(2)+1)=sum(dynamics_cell{c_escalones}(buenos));
    end % c_casos
end % c_escalones
c_bichos=N_tot+1;
estados_act = c_bichos:N_tot:c_bichos + N_tot*(c_bichos-1);
prob_final=dynamics(estados_act);