function [t,S,P0,lambda] = Program_4_2(N,beta,gamma,mu,a,S0,P0,lambda0,MaxTime)
%
% 
%
% PROGRAM_4_2( N, beta, gamma, mu, a, S0, P0, lamda0, MaxTime)
%      This is the MATLAB version of program 4.2 from page 123 of 
% "Modeling Infectious Disease in humans and animals" 
% by Keeling & Rohani.
% 
% It is a N strain SIR disease model, where the strains are arranged in a
% circle and each strain offers partial immunity (in terms of reduced
% transmission) to its neighbours.
%
% You may also wish to try:
% Program_4_2(4,40,9.98,0.02,0.25,[0.25 0.14 0.25 0.14],[0.016 0.55 0.016
% 0.55],[0.07 1e-12 0.07 1e-12],200);
%

% Sets up default parameters if necessary.
if nargin == 0
   N=4;
   beta=40.0;
   gamma=9.98;
   mu=0.02;
   a=0.4;
   S0=[0.08 0.1 0.1 0.11];
   P0=[0.4 0.3 0.3 0.29];
   lambda0=[0.15 0.02 0.1 0.01];
   MaxTime=200;
end

% Check sizes for terms that are 'vectors'
S0=CheckSize(S0,N,1,'Susceptibles');
P0=CheckSize(P0,N,1,'Partials');
lambda0=CheckSize(lambda0,N,1,'Lambda');

% Checks all the parameters are valid
CheckGreaterOrEqual(beta,0,'beta');
CheckGreaterOrEqual(gamma,0,'gamma');
CheckGreaterOrEqual(a,0,'a');
CheckGreaterOrEqual(mu,0,'m');
CheckGreaterOrEqual(S0,0,'S0');
CheckGreaterOrEqual(P0,0,'P0');
CheckGreaterOrEqual(lambda0,0,'lambda0');
CheckGreaterOrEqual(MaxTime,0,'MaxTime');

% initial conditions are set to suit default parameters
All=[S0 P0 lambda0];

% The main iteration 
options = odeset('RelTol', 1e-5);
[t, pop]=ode45(@Diff_4_2,[0 MaxTime],All,options,N,beta,gamma,mu,a);

S=pop(:,1:N); P=pop(:,N+(1:N)); lambda=pop(:,2*N+(1:N)); 

% plots the graphs with scaled colours
subplot(3,1,1)
h=plot(t,S,'-'); 
xlabel 'Time';
ylabel 'Susceptibles'

subplot(3,1,2) 
h=plot(t,P,'-');
xlabel 'Time';
ylabel 'Partially Susceptible'

subplot(3,1,3) 
h=plot(t,lambda,'-');
xlabel 'Time';
ylabel 'Force of Infection'


% Calculates the differential rates used in the integration.
function dPop=Diff_4_2(t, pop, N, beta,gamma,mu,a)

S=pop(1:N); P=pop(N+(1:N)); lambda=pop(2*N+(1:N)); 

dPop=zeros(3*N,1);

for i=1:N
    r=mod(i,N)+1;
    l=mod(i+N-2,N)+1;
    
    dPop(i) = mu - S(i)*(lambda(i)+lambda(l)+lambda(r)) - mu*S(i);
    
    dPop(i+N) = S(i)*(lambda(l)+lambda(r))  - P(i)*lambda(i) - mu*P(i);
    
    dPop(i+2*N) = beta*(S(i)+a*P(i))*lambda(i) - gamma*lambda(i) - mu*lambda(i);
end

% Does a simple check on the size and possible transpose
function Parameter=CheckSize(Parameter, L, W, str)

[l w]=size(Parameter);

if(l==W & w==L)
    Parameter=Parameter'; [l w]=size(Parameter);
end

if(l==1 & w==1)
    warning('Parameter %s is a scaler value, expanding to size %d x %d',str,L,W);
    Parameter=Parameter*ones(L,W);
    [l w]=size(Parameter);
end

if(l~=L | w~=W)
    error('Parameter %s is of size %d x %d and not %d x %d',str,l,w,L,W);
end


% Does a simple check on the value
function []=CheckGreaterOrEqual(Parameter, Value, str)

m=find(Parameter<Value);
if length(m)>0
    error('Parameter %s(%g) (=%g) is less than %g',str,m(1),Parameter(m(1)),Value);
end