import numpy as np

import matplotlib.pyplot as plt

from mpl_toolkits import mplot3d

def sus(s, i, r, d):

return -b*s*i

def inf(s, i, r, d):

return b*s*i-k*i-dd*i

def rec(s, i, r, d):

return k*i

def ded(s, i, r, d):

return dd*i

start = 0
end = 200
number = 1000 #number of steps
b = 0.3 #number of contacts per day
k = 0.005 #recovery rate per day
dd = 0.00055 #death rate per day
pop = 8000000000
S = pop #initial conditions
I = 10
R = 0
D = 0
s = S/pop
i = I/pop
r = R/pop
d = D/pop
vec = [s, i, r, d]
h = (end-start)/number #step size
t = np.arange(start, end, h)
li = []

for i in t: #using RK2 to calculate and then plot the solution

si, ii, ri, di = vec[0], vec[1], vec[2], vec[3]
sk1 = h*sus(si, ii, ri, di)
ik1 = h*inf(si, ii, ri, di)
rk1 = h*rec(si, ii, ri, di)
dk1 = h*ded(si, ii, ri, di)
sk2 = h*sus(si + 0.5*sk1, ii + 0.5*ik1, ri, di)
ik2 = h*inf(si + 0.5*sk1, ii + 0.5*ik1, ri, di)
rk2 = h*rec(si + 0.5*sk1, ii + 0.5*ik1, ri, di)
dk2 = h*ded(si + 0.5*sk1, ii + 0.5*ik1, ri, di)
sm = si + sk2
im = ii + ik2
rm = ri + rk2
dm = di + dk2
vec = [sm, im, rm, dm]
li.append(vec)

#fig = plt.figure(figsize = (10, 10))

#ax = plt.axes(projection='3d')
#ax.grid()

susceptible = [li[i][0] for i in range(len(li))]

infected = [li[i][1] for i in range(len(li))]
recovered = [li[i][2] for i in range(len(li))]
dead = [li[i][3] for i in range(len(li))]

plt.plot(t, susceptible, label = "susceptible")

plt.plot(t, infected, label = "infected")
plt.plot(t, recovered, label = "recovered")
plt.plot(t, dead, label = "deaths")
#ax.plot3D(susceptible, infected, recovered)
#ax.set_xlabel('susceptible', labelpad=20)
#ax.set_ylabel('infected', labelpad=20)
#ax.set_zlabel('recovered', labelpad=20)
plt.title("SIRD Model")
plt.legend()
plt.show()