#==============================================================================
import numpy                 as np
import matplotlib
import matplotlib.pyplot as plt
#==============================================================================
#            Programa sobre modelo discreto de competição entre espécies
#            -----------------------------------------------------------
n=81  # número de passos
alfa=0.0005  # taxa de transmissão da doença
beta=0.3     #  taxa de remoção de infectados
delt= 1      #  delta t 
s=np.zeros(n)
i=np.zeros(n)
r=np.zeros(n)
s[0] = 995  #  população inicial de suscetíveis
i[0] = 5    #  população inicial de infectados
r[0] = 0    #  população inicial de removidos
#
#     Evolução das populações 
#     -----------------------
for k in range(0,n-1):
    s[k+1]=s[k]-delt*alfa*s[k]*i[k]
    i[k+1]=i[k]+delt*(alfa*s[k]*i[k]-beta*i[k])
    r[k+1]=r[k]+delt*beta*i[k]
#
#      evolução das populações 
#
plt.axis([0,n+10,0,1000])
plt.plot(s,'r-')
plt.plot(i,'g-')
plt.plot(r,'b-')
plt.title('Evolução da epidemia no tempo')
plt.ylabel('número de indivíduos')
plt.xlabel('tempo')
plt.plot([n-15,n-11],[920,920],'r-')
plt.plot([n-15,n-11],[870,870],'g-')
plt.plot([n-15,n-11],[820,820],'b-')
plt.text(n-9,920,r'suscetíveis',color='r')
plt.text(n-9,870,r'infectados',color='g')
plt.text(n-9,820,r'removidos',color='b')
plt.savefig('sir_30.png')
plt.show()


#