# schoolVac.py
# As schoolNu.py, but includes vaccination.

import matplotlib.pyplot as plt
from frange import *


# Time unit: 1 h
beta = (10/4)/(40*8*24)
gamma = 3/(15*24)
nu = 1./(24*50) 
p = 0.0005   # vaccination rate 0.0005 or 0.0001

dt = 0.1       # 6 min
D = 300        # Simulate for D days
N_t = int(D*24/dt)   # Corresponding no of hours

ts = linspace(0, N_t*dt, N_t+1)
S = [0]*(N_t+1)
I = [0]*(N_t+1)
R = [0]*(N_t+1)
V = [0]*(N_t+1)

# Initial conditions
S[0] = 50
I[0] = 1
R[0] = 0
V[0] = 0
loss = 0

# Step equations forward in time
for n in range(N_t):
  S[n+1] = S[n] - dt*beta*S[n]*I[n] + dt*nu*R[n] - dt*p*S[n]
  V[n+1] = V[n] + dt*p*S[n]
  I[n+1] = I[n] + dt*beta*S[n]*I[n] - dt*gamma*I[n]
  R[n+1] = R[n] + dt*gamma*I[n] - dt*nu*R[n]
  loss = int(V[n+1] + S[n+1] + R[n+1] + I[n+1]) - \
         int(V[0] + S[0] + R[0] + I[0])

if loss > 0:
  print("Rounding loss:", loss)


plt.plot(ts, S, label="S")
plt.plot(ts, I, label="I")
plt.plot(ts, R, label="R")
plt.plot(ts, V, label="V")
plt.legend()
plt.xlabel('hours')
plt.ylabel('people')
plt.title("Flu over " + str(D) + " days; vac=" + str(p))
plt.show()