function [t,IH,IL] = Program_3_2(m,beta,gamma,n,I, MaxTime)
%
% 
%
% Program_3_2( m, beta, gamma, n, I, MaxTime)
%      This is the MATLAB version of program 3.2 from page 64 of 
% "Modeling Infectious Disease in humans and animals" 
% by Keeling & Rohani.
% 
% It is the SIS model with m different risk-groups
% Note we no-longer explicitly model the susceptible class.


% Sets up default parameters if necessary.
if nargin == 0
   m=5;
   beta=0.0016*[0 3 10 60 100]'*[0 3 10 60 100];
   gamma=0.2*[1 1 1 1 1];
   n=[0.06 0.31 0.52 0.08 0.03];
   I=[0.0 0.0 0.0 0.0 1e-5];
   MaxTime=30;
end

% Check the size of all inputs
A=size(beta);
if A(1)~=m | A(2)~=m
    error('Size of beta matrix does not correspond to number of risk groups');
end

if length(gamma)~=m
    error('Size of gamma vector does not correspond to number of risk groups');
end

if length(n)~=m
    error('Size of n vector does not correspond to number of risk groups');
end

if length(I)~=m
    error('Size of I vector does not correspond to number of risk groups');
end
gamma=reshape(gamma,1,m);
n=reshape(n,1,m);
I=reshape(I,1,m);


% Checks all the parameters are valid

for i=1:m
    if I(i)<0
        error('Initial level of infection in risk group %d (%g) is less than zero',i,I(i));
    end
    if n(i)<=0
        error('Risk group %d (%g) is less than or equal to zero',i,n(i));
    end
    if I(i)>n(i)
        error('Initial level of infection in risk group %d (%g) is great than the size of the group (%g)',i,I(i),n(i));
    end
    if gamma(i)<=0
        error('Recovery rate gamma(%d) (%g) is less than or equal to zero',i,gamma(i));
    end

    for j=1:m
        if beta(i,j)<0
            error('Transmission rate beta(%d,%d)=%g is less than zero',i,j,beta(i,j));
        end
    end
end

if abs(sum(n)-1)>0.1
    error('Sum of the risk groups (%g) does not equal one',sum(n));
end
if abs(sum(n)-1)>0.01
    warning('Sum of the risk groups (%g) does not equal one: n and I have been rescaled',sum(n));
end
n=n/sum(n); i=i/sum(n);
    
if MaxTime<=0 
    error('Maximum run time (%g) is less than or equal to zero',MaxTime);
end
    
if length(find(abs(beta-beta')>0.001))>0
    warning('Transmission rate beta is not symmetric');
end

% The main iteration 
options = odeset('RelTol', 1e-5);
[t, pop]=ode45(@Diff_3_2,[0 MaxTime],[I],options,[reshape(beta,1,m*m) gamma n m]);

% plots the graphs with scaled colours
str=['legend(h'];
for i=1:m
    str=[str ',''risk group' num2str(i) ''''];
end


subplot(2,2,1)
h=plot(t,pop,'-r');
for i=1:m
    z=(i-1)/(m-1);
    set(h(i),'LineWidth',2,'Color',[z 1-z 0]);
end
eval([str ',2);']);
xlabel 'Time';
ylabel 'Infectious, I'

subplot(2,2,2)
h=semilogy(t,pop,'-r');
for i=1:m
    z=(i-1)/(m-1);
    set(h(i),'LineWidth',2,'Color',[z 1-z 0]);
end
eval([str ',4);']);
xlabel 'Time';
ylabel 'Infectious, I'

subplot(2,2,3) 
h=plot(t,pop./(ones(length(t),1)*n),'-r'); 
for i=1:m
    z=(i-1)/(m-1);
    set(h(i),'LineWidth',2,'Color',[z 1-z 0]);
end
eval([str ',2);']);
xlabel 'Time';
ylabel 'Relative level of infection, I/n'

subplot(2,2,4) 
h=semilogy(t,pop./(ones(length(t),1)*n),'-r'); 
for i=1:m
    z=(i-1)/(m-1);
    set(h(i),'LineWidth',2,'Color',[z 1-z 0]);
end
eval([str ',4);']);
xlabel 'Time';
ylabel 'Relative level of infection, I/n'


% Calculates the differential rates used in the integration.
function dPop=Diff_3_2(t,pop, parameter)

m=parameter(end);
beta=reshape(parameter(1:(m*m)),m,m);  gamma=parameter((m*m)+[1:m])'; n=parameter((m*(m+1))+[1:m])';
I=pop; S=n-I;

dPop=(beta*I).*S - gamma.*I;




