% one-locus.m
% 1-locus 2-allele spatial population
% random mating within a neighborhood
% selection, mutation
%                                                           4/3/3           SG

clear all
close all
warning off
tic

s = 0.05;           % selecton strength
sigma = .0;
mu= 0.00001;       % mutation probability
n = 100;            % system size
T = 5000;          % duration

colormap([1 1 0; 0 1 0; 0 0 1]);  % yellow-green-blue

% FITNESSES
	% additive selection
	w_aa = 1;
	w_Aa = 1+s;
	w_AA =1+2*s;
	% underdominant selection
% 	w_aa   = 1;
% 	w_Aa  = 1-s;
% 	w_AA = 1+sigma;

% INITIAL CONDITIONS
if 0   X = floor(3*rand(n)); end      % uniform distribution AA(33%):Aa(33%):aa(33%)

if 1   X = zeros(n);  end                 % sea of yellow

if 0
	r = 2;                           % radius of initial pocket (blue pocket in a see of yellow)
	X = zeros(n);
	X(n/2:n/2+r,n/2:n/2+r) = 1;
end

image(X,'EraseMode','background')
set(gca,'DrawMode','fast','XTick',[],'YTick',[])
title('t = 0')
axis square

%M(1) = getframe;   % for movies
XX = ones(n);

for t=1:T

	A = X/2;                    % frequencies of A in residents
	
	Z1=X(:,[n 1:n-1]);        % states of neighbors left
	Z2=X(:,[2:n 1]);            % states of neighbors right
	Z3=X([n 1:n-1],:);        % states of neighbors above
	Z4=X([2:n 1],:);            % states of neighbors below
	B = (X+Z1+Z2+Z3+Z4)/10;             % frequency of A in neighborhood
	
	wmean = w_aa*(1-A).*(1-B)+w_Aa*(A.*(1-B)+(1-A).*B)+w_AA*A.*B; % mean fitness
	
	P_aa = w_aa*(1-A).*(1-B)./wmean;    % frequency of aa after selection
	P_AA = w_AA*A.*B./wmean;            % frequency of AA after selection
	
	Y1=rand(n);
	Y2=rand(n);
	
	X = XX;                                                 % sites with Aa
	X(Y1<P_aa) = 0;                                     % sites with aa
	X((Y1>P_aa)&(Y2<P_AA./(1-P_aa))) = 2;   % sitex with AA
	
	if mu
        Q = rand(n) < 2*mu;      % sites where mutation occurs
        X(Q&(X==0)) = 1;          % aa mutates to Aa
        X(Q&(X==2)) = 2;          % AA mutates to Aa
        QQ = rand(n);
        X(Q&(X==1)&(QQ<0.5)) = 2;   % Aa mutates to AA
        X(Q&(X==1)&(QQ>0.5)) = 0;   % Aa mutates to aa
	end
	
    aa=sum(sum(X==0))/n^2;      % overall freq. of aa
    Aa=sum(sum(X==1))/n^2;      % same of Aa
    AA=sum(sum(X==2))/n^2;      % same of AA
    
	image(XX+X,'EraseMode','background')
	set(gca,'DrawMode','fast','XTick',[],'YTick',[])
	title(['t = ', num2str(t),'; aa= ',num2str(aa,2),', Aa= ',num2str(Aa,2),', AA= ',num2str(AA,2)])
	axis square
	pause(.01)	
	%M(t+1) = getframe;                 % for movies
end

toc

%movie(M)