%Start by looking at the function called WrightFisher at the end of the
%script to understand what is done.
%Be careful if you loop over this function, avoid big numbers as
%there are already 2 loops in it.
%Because there are loops using the function it is not possible to run the
%script. You need to run the sections separately. (You can use crtl+enter
%or the "Run Section" button to do so)

%% example of simulation
close all; clear all;
N = 10; n = 100; M = 100; b = 1;
B = RunWrightFisher(N,n,M,b);

% plot 7 random trajectories
figure
for i=1:21
    plot([0:n],N.*B(i,:),'LineWidth',5)
    hold on
end
ylabel('number of individuals of type B')
xlabel('generation')

%histograms of the mean population over the m first generations
m = 20;
H = zeros(N+1,m);
for i = 1:m
    H(:,i) = histcounts(B(:,i),[-1/(2*N):1/N:1+1/(2*N)])./M;
end
figure
bar3(H)
xlabel('generation')
ylabel('individuals of type B')
set(gca,'XTickLabel',[0:m])
set(gca,'YTickLabel',[0:10])

%% fixation dependence on the population size (N)
close all; clear all;
n = 100; M = 1000; b = 1;
N = [10:10:100];
Fix = zeros(1,length(N)); %proportion of repeats with fixation of B

for i = 1:length(N)
    B = RunWrightFisher(N(i),n,M,b);
    Fix(i) = sum(B(:,end)==1)/M;
end
figure
plot(N,Fix,'LineWidth',5)
grid on
xlabel('population size')
ylabel('rate of fixation')

%% fixation dependence on the number of generations(n)
close all; clear all;
N = 10; M = 1000; b = 1;
n = [10:10:100];
Fix = zeros(1,length(n)); %proportion of repeats with fixation of B
for i = 1:length(n)
    B = RunWrightFisher(N,n(i),M,b);
    Fix(i) = sum(B(:,end)==1)/M;
end
figure
plot(n,Fix,'LineWidth',5)
grid on
xlabel('number of generations')
ylabel('rate of fixation')

%% Fixation time distribution
N = 10; n = 100; M = 5000; b = 1;
B = RunWrightFisher(N,n,M,b);
FixTime = [];
for i = 1:M
    ft = min(find(B(i,:)==1));
    if ~isempty(ft)
        FixTime = [FixTime,ft];
    end
end
figure
histogram(FixTime,'Normalization','pdf')
xlabel('generation')
title('Fixation time pdf')

%% the function simulting the Wright-Fisher model
function B = RunWrightFisher(N,n,M,b)
% N = population size
% n = number of generations
% M = number of repeats
% b = initial number of individuals of type B
% B is a (M x n+1) matrix

B = zeros(M,n+1); % number of individuals of type B for the M repeats at each generation
B(:,1) = b; %initial value

for i = 1:M
    for j = 1:n
        p = B(i,j)/N;
        B(i,j+1) = binornd(N,p); %draw the number of individuals of type B according to binomial distibution;
    end
end
B = B./N; % to have the proportion of B instead of the number of individuals
end