# -*- coding: utf-8 -*-
import numpy as np
import pylab as py



n = 100#rozmiar sieci n na n
# stan 0 ozancza podatnego
a = 8#5 #ilosc stanow infekcyjnych: 1,...,a - np. w konkretnym dniu infekcji
b = 14 #ilosc stanow z odpornoscia: a+1,...,a+b
T = 300
p = 0.1 #poczatkowa czestosc wystepowania stanow infekcyjnych i opornych


# Stan poczatkowy
D1 = np.ceil( (a+b)*np.random.rand(n,n) ) # macierz wypełniona liczbami z rozkładu płaskiego od 1 do a+b
D2 = np.random.rand(n,n) < p # wybieramy wezly  zarazeni i odporni
D = D1*D2

# PROGRAM
py.figure(1)
py.ion()
aa = py.imshow(D, interpolation = 'nearest')# stan poczatkowy
py.colorbar()
#; colormap(X);  colorbar ;
S = np.zeros(T)
I = np.zeros(T)
R = np.zeros(T)

t = 0
S[t] = np.sum(np.sum(D==0)) # poczatkowa ilosc podatnych
I[t] = np.sum(np.sum((D>0)&(D<a))) # zarazonych  (D>0)&(D<b))
R[t] = np.sum(np.sum(D>a)) #odpornych



for t in range(1,T):
    D1 = D[:,np.hstack((n-1,range(0,n-1)))] # stan sasiadow w lewo (dla każdego wiersza, 
                        # dla kolumny zero sąsiadem lewym jest n-ty, 
                        # dla pierwszej zerowy itd )
    D2 = D[:,np.hstack((range(1,n),0))] # stan sasiadow w prawo
    D3 = D[np.hstack((n-1,range(0,n-1))),:]# stan sasiadow powyzej
    D4 = D[np.hstack((range(1,n),0)),:]#stan sasiadow ponizej
    # liczymy chorych sąsiadów:
    sick_neighbors = np.where((D1>0)&(D1<a),1,0) +\
                     np.where((D2>0)&(D2<a),1,0) +\
                     np.where((D3>0)&(D3<a),1,0) +\
                     np.where((D4>0)&(D4<a),1,0)
    # zbieramy indeksy węzłów, w których trzeba coś zrobić:
    stay_healthy = ((D==0)&(sick_neighbors==0)) # podatni ktorzy nie zachoruja - tym nic nie trzeba robić
    get_sick = ((D==0)&(sick_neighbors>0)) # podatni którzy mają chorego co najmniej jednego sąsiada zachoruja
    get_recovered = (D>=(a+b)) # odporni ktorzy stana sie podatni
    others = (D>0)&(D<a+b) #wszyscy pozostali, czyli na jakimś etapie choroby lub uzyskanej czasowek odporności

    D[get_sick] = 1 
    D[get_recovered] = 0
    D[others] = D[others]+1 
    # rysujemy stan aktualny
    
    aa.set_data(D)#; colorbar; title(['t = ', num2str(t)])
    #py.draw()
    py.pause(0.1)
    # zliczamy statystyki:
    S[t] = np.sum(np.sum(D==0))
    I[t] = np.sum(np.sum((D>0)&(D<a)))
    R[t] = np.sum(np.sum(D>a))
py.ioff()  

py.figure(2)   # S,I i R
t= range(T)
py.plot(t,S,t,I,t,R)
py.legend(('podatni','zakazeni','odporni'))
py.xlabel('czas')

py.show()