% code from the web site of Sergey Gavrilets




% sir.m  -  two-dimnesional CA for SIR model
%                               SG   Feb.13, 2003                   

clear all
close all

%build the GUI
%define the plot button
plotbutton=uicontrol('style','pushbutton',...
   'string','Play', ...
   'fontsize',12, ...
   'position',[100,0,50,20], ...
   'callback', 'run=1;');

%define the stop button
erasebutton=uicontrol('style','pushbutton',...
   'string','Stop', ...
   'fontsize',12, ...
   'position',[200,0,50,20], ...
   'callback','freeze=1;');

%define the Quit button
quitbutton=uicontrol('style','pushbutton',...
   'string','Quit', ...
   'fontsize',12, ...
   'position',[300,0,50,20], ...
   'callback','stop=1;close;');


% PARAMETERS

n = 100;          % size n-by-n
a = 8;          % number of infectious states: 1,...,a
b = 8;          % number of immune states: a+1,...,a+b

p = 0.1;       % initial frequency of infectious + immune

%CREATING COLORMAP

X0 = [0 1 0];               % susceptible rae green
Xa = colormap(autumn(a));   % infected are red to yellow
Xb = colormap(winter(b));   % immune are blue to green
X= cat(1,X0,cat(1,Xa,Xb));   % overall colormap

% INITIAL STATE

%D = round((a+b)*rand(n,n));

D1 = 1+round((a+b-2)*rand(n,n));     % uniform from 1 to a+b
D2 = rand(n,n) < p;                  % those that will be infectious or immune
D = D1.*D2;

% PROGRAM

image(D);                           % initial figure
colormap(X);
axis square;
%colorbar

t = 1;
S(t) = sum(sum(D==0));              % initial number of susceptible
I(t) = sum(sum((D>0)&(D<b)));       % initial number of infectious
R(t) = sum(sum(D>a));               % initial number if immune




stop= 0;                            %wait for a quit button push
run = 0;                            %wait for a draw 
freeze = 0;                         %wait for a freeze
while (stop==0) 
    if (run==1)

		D1 = D(:,[n 1:n-1]);        % states of neighbors left
		D2 = D(:,[2:n 1]);          % states of neighbors right   
		D3 = D([n 1:n-1],:);        % states of neighbors above
		D4 = D([2:n 1],:);          % states of neighbors below
		
		sick_neighbors = ((D1>0)&(D1<b))+((D2>0)&(D2<b)) + ((D3>0)&(D3<b)) + ((D4>0)&(D4<b));
		
		stay_healthy = ((D==0)&(sick_neighbors==0)); % identify susceptibles who don't get sick
		
		get_sick = ((D==0)&(sick_neighbors>0));      % identify susceptibles who do get sick
		
		get_recovered = (D==(a+b));                  % identify immune who become susceptible
		others = (D>0)&(D<a+b);                      % identify all others
		
		D(get_sick) = 1;                % update state of susceptible who became sick
		D(get_recovered) = 0;           % update state of immune who become susceptible
		D(others) = D(others)+1;        % updata state of all other types
		
        
        % DO COLOR FIGURE
		image(D)
		title(['t = ', num2str(t)])
        colormap(X)
		%colormap(flipud(hot(a+b+1)));
		%colorbar
		axis square;
		pause(.01)
		
		t = t+1;                            % update time
            
		S(t) = sum(sum(D==0));              % overall number of susceptible
		I(t) = sum(sum((D>0)&(D<b)));       % overall number of infectious
		R(t) = sum(sum(D>a));               % overall number if immune

   
    end
    if (freeze==1)
        run = 0;
        freeze = 0;
    end
    drawnow  %need this in the loop for controls to work  
end
	


% Plot the dynamics of S,I and R

figure
T=1:t;
plot(T,S)
plot(T,S,T,R)
plot(T,S,T,R,T,I)
legend('susceptible','infectious','immune')
xlabel('time')