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


def param(n):
    X = np.linspace(-np.pi,np.pi,1000)
    rho = lambda t : 1 + np.cos(n*t) + (np.sin(n*t))**2
    plt.plot([rho(t)*np.cos(t) for t in X],[rho(t)*np.sin(t) for t in X],label = str(n))
    plt.axis('equal')
    plt.legend()
    

##plt.figure('Param')
##for i in range(7):
##    param(i)
##plt.show()
a = 1
f = lambda alpha,t : -alpha + a * np.exp(-t)

def Euler(f, alpha0, t0, dt, n):
    Lalpha = [alpha0]
    Lt = [t0]
    alpha = alpha0
    t = t0
    tmax = t0 + n*dt
    while Lt[-1] <= tmax:
        alpha = alpha + dt*f(alpha,t)
        t = t+dt
        Lt.append(t)
        Lalpha.append(alpha)
    return Lt,Lalpha

alphaexact = lambda t : np.exp(-t)
##X,Y = Euler(f,0.8,0,10**(-1),100)
##
##Y2 = odeint(f,0.8,X)
##Yr = [alphaexact(t) for t in X]
##
##
##plt.plot(X,Y,label = 'Euler')
###plt.plot(X,Yr,label = 'solution')
##plt.plot(X,Y2)
##plt.show()

S0 = 762
I0 = 1
r = 2.18 * 10**(-3)
a = 0.44
def f2(S,I,t):
    return(-r*S*I,r*S*I-a*I)

def Euler2(f,S0,I0,t0,dt,n):
    Lt,LS,LI = [t0],[S0],[I0]
    t,S,I,tmax = t0,S0,I0,t0+n*dt
    while t<= tmax:
        dS,dI = f(S,I,t)
        S,I,t = S+dt*dS,I+dI*dt,t+dt
        LS.append(S)
        LI.append(I)
        Lt.append(t)
    return(Lt,LS,LI)


def F(X,t):
    S,I = X
    return(np.array([[-r*S*I],[r*S*I-a*I]]))

def EulerVect(F,X0,t0,dt,n):
    LX = [X0]
    Lt = [t0]
    X = X0
    t = t0
    tmax = t0 + n*dt
    while Lt[-1] <= tmax:
        X = X + np.dot(dt,F(X,t))
        t = t+dt
        Lt.append(t)
        LX.append(X)
    return Lt,LX

##X,Y,Z = Euler2(f2,S0,I0,0,1/24/60,24*30*60)
##plt.figure()
##plt.plot(X,Y,label = 'Susceptible')
##plt.plot(X,Z,label = 'Infecté')
##plt.legend()
##plt.show()

X0 = (S0,I0)
T,LX = EulerVect(F,X0,0,1/24/60,24*5*60)
plt.figure()
plt.plot(T,Y,label = 'Susceptible')
plt.plot(T,Z,label = 'Infecté')
plt.legend()
plt.show()