function f = gendrift(f0,n,g,cstr)

% GENDRIFT.M: simulates genetic drift
%
% f = gendrift(f0,n,g)
%
% Inputs:
% - f0, initial allele frequency (default = 0.5)
% - n, population size (default = 1000)
% - g, number of generations
% - cstr, color to use for plot
%
% Outputs:
% - f, evolution of the allele's frequency due purely to genetic drift as a
% vector of length, g
%
% The basic idea is that we're randomly sampling from each generation
% to generate the next.  The heart of the algorithm is the discrete
% uniform distribution (unidrnd), which is used to generate a random
% set of indices for selecting the next generation's alleles.
%
% RTB wrote it, 14 July 2003

% set up some defaults
if nargin < 4, cstr = 'b'; end
if nargin < 3, g=1000; end
if nargin < 2, n=1000; end
if nargin < 1, f0 = 0.5; end

f = zeros(g,1);     % vector to hold history of allele frequencies
f(1) = f0;          % initial allele frequency

% Now we generate our 'population' which consists of a bunch of alleles.
% We'll use 0s for one allele and 1s for the other so that the frequency of
% the 1s allele is simply computed as sum(A) divided by n.
A = [ones(round(f0*n),1);zeros((n-round(f0*n)),1)]; % initial population

for k = 2:g
    % We generate the next generation by randomly drawing alleles from the
    % previous one. Remember that an allele does not get "used up" when it
    % is drawn, so we want to sample *with replacement*. This is done using
    % the discrete uniform distribution, which returns random whole numbers
    % uniformly from the set {1, 2, 3, . . . , N}. So, for example, 10
    % draws from the interval {1,5} might look like, unidrnd(5,1,10):
    % 2     1     2     5     4     2     5     1     1     1
    % The trick is to use these values as indices into the vector
    % representing the previous generation.
    A = A(unidrnd(n,n,1));  % new population is a random index into old
    f(k) = sum(A) ./ n;
end

% Now let's display our results.
plot(f,cstr);
axis([1 g 0 1]);
hl = line([1 g],[f0 f0]);
set(hl,'Color','r','LineStyle','--');
xlabel('generation #');
ylabel('allelic frequency');
tstr = sprintf('Population Size: %d', n); title(tstr);
return;

% For-loop exercise:
% Write a script that will superimpose plots of 100 simulations each with
% populations sizes of 20, 200 and 2000. Start with an allele frequency of
% 0.5 and simulate 50 generations. Each group of 10 simulations should be
% plotted in a different color.

pops = [20 200 2000];
cstrs = ['g' 'b' 'k'];
nsims = 10;
g = 50; p = 0.5;
hold on;
for k = 1:length(pops)
    for m = 1:nsims
        f = gendrift(p,pops(k),g,cstrs(k));
        drawnow;
        pause(0.2);
    end
end