# rumor.py
# Section 4.4.3 of Engel's
# "Elementary Mathematics from an Algorithmic Standpoint", p.155
# Example 3
# inputs: 100 1e-4   1000 1e-4   10 1e-3

# info:
# https://www.pythonbid.com/post/pandemic-modeling-with-python-the-sir-model
# https://kavinduhettiarachchi.medium.com/modeling-a-spread-of-an-infectious-disease-using-python-adcdd762a789

import math
import matplotlib.pyplot as plt 

NUM_ITERS = 2000

n,b = map(float, input("n b=? ").split())
   # n = no. of people; b = infection rate

xs = [0]*NUM_ITERS  # no. susceptible
xs[0] = n-1 
ys = [0]*NUM_ITERS  # no. infected
ys[0] = 1
zs = [0]*NUM_ITERS  # no recovered
 
t = 0
while t < NUM_ITERS-1:
  xs[t+1] = xs[t] - b*xs[t]*ys[t]
  ys[t+1] = ys[t] - b*ys[t]*(n-2*xs[t])
  zs[t+1] = zs[t] + b*ys[t]*(n-xs[t])
  t += 1

ts = [i for i in range(NUM_ITERS)]  # time steps

print("Final set sizes")
print(f"No. susceptible (x): {xs[NUM_ITERS-1]:.1f}")
print(f"No. infected (y): {ys[NUM_ITERS-1]:.1f}")
print(f"No. recovered (z): {zs[NUM_ITERS-1]:.1f}")

plt.plot(ts, xs, label="Susceptible (X)", ls="--") 
plt.plot(ts, ys, label="Infected (Y)") 
plt.plot(ts, zs, label="Recovered (Z)", ls=":") 
plt.title("Rumors Spread (n="+str(n)+", b=" + str(b)+")")
plt.legend()
plt.xlabel("Time")
plt.ylabel("People")
plt.show()