close all;
clear all;
N = 100;

infectionRate = 0.4;
healRate = 0.05;
dieAfter = 10;
latentUntil = 4;



initialInfection = 0.001;
pop = zeros(N+2);
pop(2:N+1,2:N+1) = ones(N) + (rand(N) < initialInfection);
%pop(2:N+1,2:N+1) = ones(N);
pop(2,2) = 2;
pop(1,1) = dieAfter+2;
newpop  = pop;

map = zeros(dieAfter+2,3);
map(1,:) = [0 1 0];
map(2,:) = [1 1 1];
map(3:end, 1) = flipud(linspace(1,0,size(map(3:end, 1),1)));

for t= 1:50
    imagesc(pop);
    colormap(map);
    for i = 2:N+1
        for j = 2:N+1
            guy = pop(i,j);
            
            %immunized
            if guy == 0
                continue;
            %dead, proceed to next
            elseif guy >= dieAfter
                continue;
            %latent, infects, but doesn't heal
            elseif guy > 1 && guy < latentUntil
                newpop(i,j) = guy+1;
            %sick, if lucky get healed and immunized
            elseif guy >= latentUntil
                if rand < healRate
                    newpop(i,j) = 0;
                else
                    newpop(i,j) = guy+1;
                end
            %healthy, if anyone is sick, get infected
            elseif guy == 1
                newpop(i,j) = 1 + (sum(sum(rand(3)<infectionRate & (pop(i-1:i+1, j-1:j+1) > 1) & (pop(i-1:i+1, j-1:j+1) < dieAfter))) >= 1);
            end    
        end
    end
    pause(0.2);
    pop = newpop;
    
end