import numpy as np
import matplotlib.pyplot as plt
import functools
import matplotlib.animation as animation

swiss = np.array([2.100350058, 3.150525088, 4.900816803, 6.534422404, 10.501750292, 13.302217036, 24.970828471, 31.271878646, 39.323220537, 43.640606768, 57.292882147, 76.079346558, 100.116686114, 131.271878646, 158.576429405, 256.709451575, 256.709451575, 268.378063011, 309.218203034, 353.325554259, 453.675612602])
swiss = np.array([18, 27, 42, 56, 90, 114, 214, 268, 337, 374, 491, 652, 858, 1125, 1359, 2200, 2200, 2300, 2650, 3028, 3888])

days = np.array(range(1,len(swiss)+1))

def update(i, lines, t, p):
    n_pop, length = p.shape
    for j in range(0, n_pop):
        lines[j].set_data(t[0:i], p[j][0:i])
    return lines

# def policy_r0(R_0, R_target, t, delta_t):
#     if (t <= )



def seir(y, t, R_0, Tinf, Tinc):
    N = np.sum(y)
    y1 = np.zeros(4)
    y1[0] = - R_0 / Tinf * y[2] * y[0] / N
    y1[1] =  R_0 / Tinf * y[2] * y[0] / N - 1.0 / Tinc * y[1]
    y1[2] = 1.0 / Tinc * y[1] - 1.0 / Tinf * y[2]
    y1[3] = 1.0 / Tinf * y[2]

    return y1

def seihse(y, t, R_0, Tinf, Tinc, Thos, Tsev):
    N = np.sum(y)
    y1 = np.zeros(6)
    y1[0] = - R_0 / Tinf * y[2] * y[0] / N # Sane
    y1[1] =  R_0 / Tinf * y[2] * y[0] / N - 1.0 / Tinc * y[1] # Exposed
    y1[2] = 1.0 / Tinc * y[1] - 1.0 / Tinf * y[2] - 1.0 / Tinf * y[2] # Infectious
    y1[3] = 1.0 / Tinf * y[2] + 1.0 / Thos * y[4] + 1.0 / Tsev * y[5] # Recovered
    y1[4] = - 1.0 / Thos * y[4] + 1.0 / Tinf * y[2] - 1.0 / Tsev * y[4] # Hospitalized
    y1[5] = - 1.0 / Tsev * y[5] + 1.0 / Thos * y[4] # Severe

    return y1
    
def rk4(F, y, t, dt):
    k1 = dt * F(y, t)
    k2 = dt * F(y + k1 / 2, t + dt/2)
    k3 = dt * F(y + k2 / 2, t + dt/2)
    k4 = dt * F(y + k3, t + dt)
    return y + (k1 + 2 * (k2 + k3) + k4) / 6

R_0 = 4.26
R_01 = 4.26
R_02 = 1 - 1/R_01

Tinf = 7.0
Tinc = 5.1
Thos = 14.0
Tsev = 14.0

N = 500000

I0 = 214 / 0.2
E0 = 2000
R0 = 0
S0 = N-E0-I0
H0 = 0
Sev0 = 0
t0 = 24

# max_t = 5*days[len(swiss)-1]
max_t = 200.0
n_steps = 1000
dt = max_t / n_steps

y0 = np.array([S0, E0, I0, R0, H0, Sev0])

y_list = [y0]
t_list = [t0]
for i in range(0, n_steps):
    t = t_list[i] + dt
    # foo = functools.partial(seir, R_0=R_0, Tinf=Tinf, Tinc=Tinc)
    foo = functools.partial(seihse, R_0=R_0, Tinf=Tinf, Tinc=Tinc, Thos=Thos, Tsev=Tsev)
    y1 = rk4(foo, y_list[i], t, dt)
    y_list.append(y1)
    t_list.append(t)

    if (y1[1] + y1[2] < 1):
        break

t = np.array(t_list) 
y = np.array(y_list) 

p = np.transpose(y)


fig = plt.figure()
# ax = plt.axes(xlim=(-0.1*np.max(t),np.max(t)*1.1), ylim=(-0.1*np.max(p),np.max(p)*1.1))
ax = plt.axes(xlim=(1,np.max(t)*1.1), ylim=(1,np.max(p)*1.1))

lines = [plt.semilogy([], [], 'r-')[0],
         plt.semilogy([], [], 'b-')[0],
         plt.semilogy([], [], 'g-')[0],
         plt.semilogy([], [], 'k-')[0],
         plt.semilogy([], [], 'c-')[0],
         plt.semilogy([], [], 'y-')[0] ]

# anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), 
#            frames=n_steps, interval=10, blit=True)
ax.legend(['Sain', 'Exposé', 'Infectieux', 'Rétablis', 'Hospitalisés', 'Soins intensifs'])

anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), 
           frames=n_steps, interval=10, blit=True)

# plt.semilogy(t, p[0], 'b')
# plt.semilogy(t, p[1], 'r')
# plt.semilogy(t, p[2], 'k')
# plt.semilogy(t, p[3], 'g')
# # plt.semilogy(days, swiss, 'k*')
# plt.legend(['S', 'E', 'I', 'R', 'swiss'])
plt.show()