import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt


def sir(beta, t):
    # Total population, N.
    N = 1000
    # Initial number of infected and recovered individuals, I0 and R0.
    I0, R0 = 1, 0
    # Everyone else, S0, is susceptible to infection initially.
    S0 = N - I0 - R0
    # Contact rate, beta, and mean recovery rate, gamma, (in 1/days).
    gamma = 1.0 / 10

    # The SIR model differential equations.
    def deriv(y, t, N, beta, gamma):
        S, I, R = y
        dSdt = -beta * S * I / N
        dIdt = beta * S * I / N - gamma * I
        dRdt = gamma * I
        return dSdt, dIdt, dRdt

    # Initial conditions vector
    y0 = S0, I0, R0
    # Integrate the SIR equations over the time grid, t.
    ret = odeint(deriv, y0, t, args=(N, beta, gamma))
    S, I, R = ret.T
    return I


# A grid of time points (in days)
t = np.linspace(0, 160, 160)
I = sir(0.8, t)

# Plot the data on three separate curves for S(t), I(t) and R(t)
fig = plt.figure(facecolor="w")
ax = fig.add_subplot(111, axisbelow=True)
ax.set_facecolor = "#dddddd"
ax.plot(t, I / 1000, "r", alpha=0.5, lw=2, label="sin medidas de protección")
ax.plot(t)
ax.axhline(0.2, label="capacidad del sistema sanitario", color="g")
ax.set_xlabel("tiempo en días")
ax.set_ylabel("infectados en miles")
ax.set_ylim(0, 1.2)
ax.yaxis.set_tick_params(length=0)
ax.xaxis.set_tick_params(length=0)
ax.grid(b=True, which="major", c="w", lw=2, ls="-")
legend = ax.legend()
legend.get_frame().set_alpha(0.5)
for spine in ("top", "right", "bottom", "left"):
    ax.spines[spine].set_visible(False)
plt.show()