# Mini 2 Mand D
import numpy as np
from epidemics import *
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
from matplotlib import cm

# Settings
T = 500 #number of weeks
alpha = 0.30    # high alpha value: easy to infect
beta = 0.05 # high beta-value: easy recovering

# Initialize the grid
grid = np.loadtxt('worldmap.dat', dtype=int, delimiter=',')
grid[30:35, 265:270] = INFECTED # Start in Asia
grid[30:35, 50:55] = INFECTED # Start in Northa America

grids = []
grids.append(grid)
S = [np.sum(grid == SUSCEPTIBLE)]
I = [np.sum(grid == INFECTED)]
R = [np.sum(grid == RECOVERED)]

# Run the simulation
for n in range(T):
    grid = time_step(grid, alpha, beta)
    grids.append(grid)
    S.append(np.sum(grid == SUSCEPTIBLE))
    I.append(np.sum(grid == INFECTED))
    R.append(np.sum(grid == RECOVERED))

# Plot the results
[plot2D_SIR(grids[t], f'Week {t}') for t in np.arange(0,T+1,T//5)]

#plot evolution
fig, ax = plt.subplots()
people = [S, I, R]
lab = ['Susceptible', 'Infectious', 'Recovered']
col = ['indianred', 'darkslategrey', 'darkorange']
for i in range(len(people)):
    ax.plot(range(T+1), people[i], color=col[i], label=lab[i])
ax.legend()
plt.show()