# @author:                hehaoran
# @environment: python 3.7.4
# pylint: disable = no-member
import numpy as np
import pandas as pd
import scipy.integrate as spi

# alpha: the immunity probability.
# beta: the likelihood that the disease will be transmitted from an infected to a susceptible individual in a unit time is β.
# gamma: the recovery rate and in SI model, gamma = 0.
# I0: is the initial fraction of infected individuals, and 0 < I0 < 1.

def SI_SIS(X, beta, gamma=0):
    Y = np.zeros(2)
    Y[0] = - beta * X[0] * X[1] + gamma * X[1]
    Y[1] = beta * X[0] * X[1] - gamma * X[1]
    return Y

def SIR(X, beta, gamma):
    Y = np.zeros(3)
    Y[0] = - beta * X[0] * X[1]
    Y[1] = beta * X[0] * X[1] - gamma * X[1]
    Y[2] = gamma * X[1]
    return Y

def SIRS(X, alpha, beta, gamma):
    Y = np.zeros(3)
    Y[0] = - beta * X[0] * X[1] + alpha * X[2]
    Y[1] = beta * X[0] * X[1] - gamma * X[1]
    Y[2] = gamma * X[1] - alpha * X[2]
    return Y

def SEIR(X, alpha, beta, gamma1, gamma2):
    Y = np.zeros(4)
    Y[0] = - beta * X[0] * X[2]
    Y[1] = beta * X[0] * X[2] - (alpha + gamma1) * X[1]
    Y[2] = alpha * X[1] - gamma2 * X[2]
    Y[3] = gamma1 * X[1] + gamma2 * X[2]
    return Y

def epidemic_estimate(model, y0, T):
    """
    y0 : array-like
        the initial infected rate. For example, SI model, np.array([0.1,0.9])
    """
    return spi.odeint(model, y0, T)

if __name__ == "__main__":
    pass