"""
Code from lecture 14.11.2022. Example class for SEIR
model from lecture notes, including a test function
to test the call method.
"""

import numpy as np
import matplotlib.pyplot as plt
from ODESolver import *

class SEIR:
    def __init__(self, beta, mu, nu, gamma):
        self.beta = beta
        self.mu = mu
        self.nu = nu
        self.gamma = gamma

    def __call__(self,u,t):
        S, E, I, R = u
        N = S+I+R+E
        beta = self.beta
        mu = self.mu
        nu = self.nu
        gamma = self.gamma
        dS = -beta*S*I/N + gamma*R
        dE  = beta*S*I/N - mu*E
        dI =  mu*E - nu*I
        dR = nu*I -gamma * R
        return [dS,dE,dI,dR]


def test_SEIR():
    beta = 1.0; mu=1.0/5
    nu = 1.0/7; gamma = 1.0/50
    tol = 1e-8
    test_problem = SEIR(beta=beta,mu=mu,nu=nu,gamma=gamma)

    u = [1,1,1,1]
    t = 0
    computed = test_problem(u,t)
    expected = [-beta/4+gamma,beta/4-mu,mu-nu,nu-gamma]
    #print(computed)
    #print(expected)

    for c,e in zip(computed,expected):
        msg = f'Expected {e}, but got {c}'
        success = abs(c-e) < tol
        assert success, msg


#call the test function:
test_SEIR()


#define and solve the problem:
S_0 = 5e6
E_0 = 100
I_0 = 0
R_0 = 0

beta = 1.0; mu=1.0/5
nu = 1.0/7; gamma = 1.0/50

problem = SEIR(beta,mu,nu,gamma)
U0 = [S_0, E_0, I_0, R_0]


solver = RungeKutta4(problem)
solver.set_initial_condition(U0)
time_points = np.linspace(0, 100, 101)
u, t = solver.solve(time_points)
S = u[:,0]; E = u[:,1];
I = u[:,2]; R = u[:,3];

plt.plot(t,S,label='S')
plt.plot(t,E,label='E')
plt.plot(t,I,label='I')
plt.plot(t,R,label='R')
plt.legend()
plt.savefig('seir_class0.pdf')
plt.show()

"""
Terminal> python SEIR_class.py
(output is a plot)
"""