# this python module was written to run within the SAGE environment. See the attached 
# Jupyter notebook (.ipynb) for examples. 
# Ying Lin @ SRJC, 12/06/2020

import scipy.integrate
import numpy as np
import matplotlib.pyplot as plt
from sage.plot.plot3d.list_plot3d import list_plot3d
from sage.calculus.var import *
from sage.plot.plot3d.plot_field3d import plot_vector_field3d
from sage.plot.plot3d.shapes2 import point3d
from sage.plot.point import point2d
from sage.plot.plot_field import plot_vector_field

# initial values for the ODE

S0 = 0.99 # initial percent of susceptible people
I0 = 0.01 # initial percent of infected 
R0 = 0.0 # initial percent of recovered

# initial value used only in the SEIR model
E0 = 0.01 # initial percent of exposed

# parameters
a = 0.02 # parameter controlling the rate the resistant change back to susceptible. 
b = 0.2 # parameter controlling the rate that the susceptible becomes infected
g = 0.05 # parameter controlling the rate the infected patients recover


# time vector
t = np.linspace(0, 200, 10000)

# basic model
def SIR_system(y, t, alpha, beta, gamma):
    # S: Susceptible, I: Infected, R: Recovered
    S, I, R = y
    # This defines a system of 1st order ordinary differential equations
    dS_dt = -beta * S * I 
    dI_dt = beta * S * I - gamma * I
    dR_dt = gamma * I
    return ([dS_dt, dI_dt, dR_dt])
    

def SIR_system_with_feedback(y, t, alpha, beta, gamma):
    # variation of model based on the loss of resistance
    S, I, R = y
    dS_dt = -beta * S * I + alpha * R
    dI_dt = beta * S * I - gamma * I
    dR_dt = gamma * I - alpha * R
    return ([dS_dt, dI_dt, dR_dt])

def plot_VF(alpha = a, beta=b, gamma=g, dimension=2):
    S,I,R = var('S,I,R')
    dS_dt = -beta * S * I
    dI_dt = beta * S * I - gamma * I
    dR_dt = gamma * I
    if dimension>2:
        p = plot_vector_field3d((dS_dt, dI_dt, dR_dt), (S, 0, 1), (I, 0, 1), (R, 0, 1))
    else:
        p = plot_vector_field((dS_dt, dI_dt), (S, 0, 1), (I,0,1))       
    return p
    
# basic model plus exposed 
def SEIR_system(y, t, alpha, beta, gamma):
    # S: Susceptible, E: Exposed, I: Infected, R: Recovered
    S, E, I, R = y
    # This defines a system of 1st order ordinary differential equations
    dS_dt = -beta * S * I
    dE_dt = beta * S * I - beta * E
    dI_dt = beta * E - gamma * I
    dR_dt = gamma * I
    return ([dS_dt, dE_dt, dI_dt, dR_dt])

# use this to plot the solutions as functions of time
def run():
    res = solve_SIR()
    plot_solution(res)
    
# uses scipy ODE solver to numerically solve the system    
def solve_SIR(S=S0, I=I0, R=R0, alpha=a, beta=b, gamma=g):
    # use this if you are running the basic SIR model
    res = scipy.integrate.odeint(SIR_system, [S, I, R], t, args=(alpha, beta, gamma))
    
    # this one for the SIR with feedback. The parameters are set so that the critical point is 
    # a stable spiral point. 
    # res = scipy.integrate.odeint(SIR_system_with_feedback, [S, I, R], t, args=(alpha, beta, gamma))
    
    res = np.array(res)
    return res

# plot the phase portrait together with the slope/direction field. Default to S-I plane. 3D is slow. 
def plot_solution_with_VF(S=S0, I=I0, R=I0, dimension=2):
    res = solve_SIR(S, I, R)
    if dimension == 2:
        curve = point2d(res[:,:2])
    else:
        curve = point3d(res)
    field = plot_VF(a, b, g, dimension)
    return(curve, field)

# plot the phase portrait with multiple initial values. 
def plot_orbits_IV(dimension=2):
    initialValues = [(0.9, 0.1, 0), (0.4, 0.2, 0.4), (0.6, 0.2, 0.2), (0.2, 0.1, 0.7), (0.2,0.8,0), (0.15, 0.4, 0.45), (0.3, 0.6, 0.1)]
    graphList = []
    for iv in initialValues:
        (c, f) = plot_solution_with_VF(iv[0], iv[1], iv[2], dimension)
        graphList.append(c)
    field = plot_VF(a, b, g, dimension)
    return(graphList, field)    
    
# plot S, I, R as a functions of time. 
def plot_solution(res):
    # plot
    plt.figure(figsize=[6, 4])
    plt.plot(t, res[:, 0], label='Susceptible')
    plt.plot(t, res[:, 1], label='Infected')
    plt.plot(t, res[:, 2], label='Recovered')
    plt.legend()
    plt.grid()
    plt.xlabel('time')
    plt.ylabel('proportions')
    plt.title('SIR model simulation')
    plt.show()

# plotting the SEIR model 
def plot_solution_SEIR(res):
    # plot
    plt.figure(figsize=[6, 4])
    plt.plot(t, res[:, 0], label='Susceptible')
    plt.plot(t, res[:, 1], label='Exposed')
    plt.plot(t, res[:, 2], label='Infected')
    plt.plot(t, res[:, 3], label='Recovered')
    plt.legend()
    plt.grid()
    plt.xlabel('time')
    plt.ylabel('proportions')
    plt.title('SEIR model simulation')
    plt.show()