N = 100;

dieAfter = 10;
initialInfection = 0.001;
infectionRate=0.4;
healRate=0.1;

%everyone on the boundary is immune
pop = zeros(N+2); 

%for starters, set everyone to healthy, and infect a few
pop(2:N+1,2:N+1) = ones(N) + (rand(N) < initialInfection); 

%create a temporary array
newpop  = pop;

%create colourmap
%you can be immune, healty or at different stages of the disease (up to being dead)
map = zeros(dieAfter+2,3);
map(1,:) = [0 1 0]; % zero is immune, colour green
map(2,:) = [1 1 1]; % one is healthy, colour white
map(3:end, 1) = flipud(linspace(1,0,size(map(3:end, 1),1))); % shades of red (up to black)

for t= 1:50
    %show current state
    imagesc(pop,[0 dieAfter+2]);
    colormap(map);
   for i=2:N+1
       for j=2:N+1
          guy=pop(i, j);
          newguy=newpop(i, j);
          if guy>=dieAfter %ha halott, úgy marad
              
          elseif guy==0 %ha immunis, az is marad
          elseif guy==1 %fertõzhetõ
              korny=0;
              for i1=-1:1  %megszámolom, hogy az 1 sugarú környezetében kik az élõ fertõzöttek
                  for j1=-1:1
                      if pop(i+i1, j+j1)>1 && pop(i+i1, j+j1)<dieAfter
                          korny=korny+1;
                      end
                  end
              end
              if korny*rand>infectionRate %megfertõzõdhet az õt körbevevõ fertõzöttek számától függõen
                  newguy=2;
              end
              
          elseif guy>1 % ha beteg:
              if rand<healRate % bizonyos valószínûséggel meggyógyul és immunis lesz
                  newguy=0;
              else
                  newguy=guy+1; %ha nem gyógyult meg, még betegebb lesz
              end
          end
          newpop(i,j)=newguy;
       end
   end
    
    waitforbuttonpress;
    %pause(0.2);
    %update population with the temporary array
    pop = newpop;
    
end