#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import numpy as np
import matplotlib.pyplot as plt


# In[ ]:


# S0, IWT0, R0: initial sizes of subpopulations
# a:  removal rate of infected, rWT infection rate
# T:  solve the equations until time T
# ts: array of time points

def solve_SIR(S0, IWT0, R0, rWT, a, ts):
    
    #initialize the populations vectors
    S = np.zeros(ts.size)
    IWT = np.zeros(ts.size)
    R = np.zeros(ts.size)
    N = S0 + IWT0 + R0
    S[0] = S0
    IWT[0] = IWT0
    R[0] = R0
    
    # Euler method loop
    for tidx in range(1,len(ts)):
        # update S I R using the values from the last timestep
        S[tidx] = S[tidx-1] - (rWT*IWT[tidx-1]*S[tidx-1]/N) * dt
        IWT[tidx] = IWT[tidx-1] + (rWT*IWT[tidx-1]*S[tidx-1]/N - a*IWT[tidx-1]) * dt
        R[tidx] = R[tidx-1] + (a*IWT[tidx-1]) * dt
        
    return S, IWT, R


# In[ ]:


a = 1
rWT = 3
T = 5

plt.figure()

dt = 1.0e-5
ts = np.arange(0, T, dt)
S, IWT, R = solve_SIR(8000, 50, 1950, rWT, a, ts)
plt.plot(ts, S)
plt.plot(ts, IWT)
plt.plot(ts, R)

plt.show()


# In[ ]: