% Wright-Fisher process for a costly signaling game in which there are
% two types. 

% This code calculate time trends for a population of size pop_span over 
% number of generations given by num_gens. Mutation rate is given by u_span. 
% The cost of signal to hi type and lo type, respectively, are given 
% by c_lo and c_hi. 


clc
clear all

tic

% Determine population size.
pop_size = 100;

% Determine number of generations to run.
num_gens = 10000;

wage = 5;

% Define the probability of being low type. 
q = 2/3;

% Determine number of receivers (half the population)
num_receivers = round(pop_size/2);

% Determine the type of each non-receiver player.
num_hi = round(pop_size*(1-q)/2);
num_lo = round(pop_size*(q)/2);

% Define selection strength. 
w = 1;

% Define mutation rate. 
u = 0.01;

% Set costs of signals by type.
c_lo = [0, 3, 6, 12];
c_hi = [0, 1, 2, 4];

% Store number of signals. 
[r signals] = size(c_hi);


% Define payoffs from choosing each type (first entry is payoff to receiver for 
% choosing low type, second is payoff to receiver for choosing high type. 
receiver_pi = [0, 10];

% Initialize matrix sizes (pre-allocate memory). NOTE: For each agent, 
% strats_hi and strats_lo will be of form 0 0 1 0 each period. For each
% receiver, strats_receiver will be of the form 0 1 1 1 in each period.

strats_hi = zeros(num_hi, signals, num_gens);
strats_lo = zeros(num_lo, signals, num_gens); 

payoff_hi = zeros(num_hi, 1, num_gens);
payoff_lo = zeros(num_lo, 1, num_gens);

strats_receivers = zeros(num_receivers, signals, num_gens); 
payoff_receivers = zeros(num_receivers, 1, num_gens);

avg_strat_hi = zeros(num_gens, signals);
avg_strat_lo = zeros(num_gens, signals);
avg_strat_receivers = zeros(num_gens, signals);


% Set initial values of signals being sent by lo type. 
for k=1:1:num_lo
    number = floor(rand*signals) + 1;
    strats_lo(k, number, 1) = 1;
end

% Set initial values of signals being sent by hi type. 
for k=1:1:num_hi
    number = floor(rand*signals)+1;
    strats_hi(k, number, 1) = 1;
end


strats_receivers(:,:,1) = round(rand(num_receivers,signals));

avg_strat_hi(1,:) = mean(strats_hi(:,:,1));
avg_strat_lo(1,:) = mean(strats_lo(:,:,1));
avg_strat_receivers(1,:) = mean(strats_receivers(:,:,1));

avg_cost_hi(1) = avg_strat_hi(1,:)*c_hi';
avg_cost_lo(1) = avg_strat_lo(1,:)*c_lo';


for t = 1:1:num_gens
        
    % Calculate the payoffs for each individual.
    choosing_lo = strats_lo(:,:,t)*strats_receivers(:,:,t)';
    payoff_lo(:,1,t) = (sum(choosing_lo')')*wage/num_receivers - strats_lo(:,:,t)*c_lo';
    
    choosing_hi = strats_hi(:,:,t)*strats_receivers(:,:,t)';
    payoff_hi(:,1,t) = (sum(choosing_hi')')*wage/num_receivers - strats_hi(:,:,t)*c_hi';
    
    % Calculate payoffs for each receiver.
    payoff_if_choose_hi = strats_receivers(:,:,t)*strats_hi(:,:,t)';
    num_hi_chosen = sum(payoff_if_choose_hi')'; % returns number of hi types chosen by receivers
    payoff_if_chose_lo = strats_receivers(:,:,t)*strats_lo(:,:,t)';
    num_lo_chosen = sum(payoff_if_chose_lo')'; % returns number of lo types chosen by receivers
    
    payoff_receivers(:,:,t) = receiver_pi(1)*num_lo_chosen + receiver_pi(2)*num_hi_chosen - wage*(num_lo_chosen + num_hi_chosen);
    
    % UPDATING.(1) Update all the lo types, then (2) update
    % all the hi types, then (3) update all the receivers.
    
    % (1) Factor in selection strength, w, for lo type.
    selection_lo = exp(payoff_lo(:,1,t)*w);
    
    % Find each lo type's mother.
    for i=1:1:num_lo
        
        random_selection = rand*sum(selection_lo);
        momma = 1;
        
        while sum(selection_lo(1:momma)) < random_selection
            momma = momma+1;
        end
        
        % Lo types reproduce, do we mutate?
        mutate = rand;
        
        if mutate > u
            strats_lo(i,:,t+1) = strats_lo(momma,:,t);
        else
            number = floor(rand*signals) + 1;
            strats_lo(i,number,t+1) = 1;
        end
    end
    
    % (2) Factor in selection strength, w, for hi type.
    selection_hi = exp(payoff_hi(:,1,t)*w);
    
    % Find each hi type's mother.
    for i=1:1:num_hi
        
        random_selection = rand*sum(selection_hi);
        momma = 1;
        
        while sum(selection_hi(1:momma)) < random_selection
            momma = momma+1;
        end
        
        % Reproduce, do we mutate?
        mutate = rand;
        
        if mutate > u
            strats_hi(i,:,t+1) = strats_hi(momma,:,t);
        else
            number = floor(rand*signals) + 1;
            strats_hi(i,number,t+1) = 1;
        end
    end
    
    % (3) Factor in selection strength, w, for receivers. 
    selection_receivers = exp(payoff_receivers(:,1,t)*w);
    
    % find everyone's mother
    for i=1:1:num_receivers
        
        random_selection = rand*sum(selection_receivers);
        momma = 1;
        
        while sum(selection_receivers(1:momma)) < random_selection
            momma = momma+1;
        end
        
        % Reproduce, do we mutate?
        mutate = rand;
        
        if mutate > u
            strats_receivers(i,:,t+1) = strats_receivers(momma,:,t);
        else
            strats_receivers(i,:,t+1) = round(rand(1,signals));
        end
    end
    
    % Find average signal sent for hi type. 
    avg_strat_hi(t+1,:) = mean(strats_hi(:,:,t+1));
    
    % Find average signal sent for lo type. 
    avg_strat_lo(t+1,:) = mean(strats_lo(:,:,t+1));
    
    % Find average signal sent for receivers. 
    avg_strat_receivers(t+1,:) = mean(strats_receivers(:,:,t+1));
    
    
    avg_cost_hi(t+1) = avg_strat_hi(t+1,:)*c_hi';
    avg_cost_lo(t+1) = avg_strat_lo(t+1,:)*c_lo';
    
end


figure
plot(1:length(avg_cost_hi),avg_cost_hi, 'b',...
    1:length(avg_cost_hi),avg_cost_lo, 'r')
legend('Avg cost to hi type','Avg cost to low type' )
title(strcat('Costly Signaling in Wright Fisher, mutation=', num2str(u)))
xlabel('Generation')
ylabel('Average signal sent')
hold on


toc



