%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import math

N = 1000
I = 1
S = N - I
R = 0

beta = 0.0002
gamma = 0.05

sim_length = 120
dt = 0.01

xs = []
y1s = []
y2s = []
y3s = []

num_iter = int(sim_length / dt) + 1
for i in range(num_iter):
    xs.append(i * dt)
    y1s.append(S)
    y2s.append(I)
    y3s.append(R)
    Sold = S
    Iold = I
    Rold = R
    S += (-beta * Iold * Sold) * dt
    I += (beta * Iold * Sold - gamma * Iold) * dt
    R += (gamma * Iold) * dt
    
plt.plot(xs, y1s)
plt.plot(xs, y2s)
plt.plot(xs, y3s)
plt.xlabel("Time")
plt.ylabel("Population")
plt.title("SIR Epidemiology Model")
plt.legend(['Susceptible', 'Infected', 'Recovered'])