function [nrecovered,infected,recovered,susceptible,vaccinated] = ...
             modelmc(n,k,tau,delta,nu,movieon)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% function [nrecovered,infected,recovered,susceptible,   %
%  vaccinated] = modelmc(n,k,tau,delta,nu,movieon)       %
%  for                                                   %
%  A solution to CSE Your Homework Assignment Project 5  %
%  Models of Infection: Person to Person                 %
%  Dianne P. O'Leary   09/03                             %
%                                                        %
% The parameters for the model:                          %
%                                                        %
%   n:       dimension of the grid of hospital beds      %
%   k:       number of days of contagion                 %
%   tau:     transmission rate for the infection         %
%   delta:   mobility rate of the population             %
%   nu:      vaccination rate                            %
%   movieon: 1 if movie is desired, 0 otherwise          %
%                                                        %
% We model the spread of an infection through a          %
% population, assuming that a person once infected       %
% cannot be reinfected.                                  %
%                                                        %
% Output:                                                %
%                                                        %
%  nrecovered  = number of patients eventually infected  %
%  infected(t) = number of individuals infected on day t %
%  recovered(t) = number of individuals who have         %
%                 recovered from infection on day t      %
%                 and are immune                         %
%  susceptible(t) = number of individuals susceptible to %
%                   infection on day t                   %
%  vaccinated(t) = number of individuals who have been   %
%                  vaccinated on day t.                  %
%  Note that infected(t) + recovered(t) + susceptible(t) %
%            + vaccinated(t) = n^2.                      %
%                                                        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

nn = n*n;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Pop(i,j) =  0  means not infected                      %
%          =  s  means in s-th day of infection          %
%          = -1  means recovered                         %
%          = -2  means vaccinated                        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Pop = zeros(n,n);
nPop = zeros(n,n);

Pop(n/2,n/2) = 1;
ninfected = 1;

t = 0;

while (ninfected > 0)

   t = t + 1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% For the next timestep, calculate the infection status  %
% of each individual.                                    %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

   for i=1:n,
   for j=1:n,
      new = Pop(i,j);
      if (new > 0)
         new =  new + 1;
         if (new > k)  % if past the time of active infection 
            new = -1;  % then recovered
         end
      elseif (new == 0)
         new = infect(Pop,i,j,n,n,tau);
      end
      nPop(i,j) = new;
   end
   end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Vaccinate                                              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

   if (nu > 0)
      for i=1:n,
      for j=1:n,
         if (rand < nu     ) & (nPop(i,j) == 0)
            nPop(i,j) = -2;
         end
      end
      end
   end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Calculate the new positions of the individuals.        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

   if (delta > 0)
      for i=1:n,
      for j=1:n,
         if (rand < delta)  
            inew = floor(rand*n+1);
            jnew = floor(rand*n+1);
            tt = nPop(i,j);
            nPop(i,j) = nPop(inew,jnew);
            nPop(inew,jnew) = tt;
          end
      end
      end
   end

   Pop = nPop;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Counts                                                 %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   infected(t)    = sum(sum(Pop >   0));
   recovered(t)   = sum(sum(Pop == -1));
   susceptible(t) = sum(sum(Pop ==  0));
   vaccinated(t)  = sum(sum(Pop == -2));
   ninfected = infected(t);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Plot the results, showing the infected.                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

   if (movieon)
     imagesc(Pop)
     title(sprintf('Day %d',t))
     colorbar
     drawnow
     if (infected(t) == 0)
        break
     end
   end


end

nrecovered = recovered(t);