# Module for simulating infectious disease dynamics with the SIR (Susceptible-Infected-Recovered) model

import numpy as np, scipy.integrate, scipy.optimize, sys

class SIRsystem:

    def __init__(self, beta, gamma, S, I, R):
        """Initialize an instance of an SIRsystem, providing values for the model parameters beta and gamma,
        and the initial numbers of hosts in each of the model compartments (S = Susceptible, I = Infectious,
        R = Recovered).

        Initialize internal instance variables self.t (time), self.S, self.I, self.R, self.N=S+I+R and
        a numpy array self.trajectory to hold [S,I,R] numbers at every time step.

        This is a base class for subsequent specialization, depending on whether one wants to simulate a
        deterministic or stochastic model of SIR dynamics.
        """
        pass

    def reset(self, S, I, R, t=0.):
        """Reset the system by setting S,I,R and t to specified values, and reinitializing a trajectory array with this
        starting configuration"""
        pass
        
    def get_total_infected(self):
        """Return the total number of hosts either currently or previously infected"""
        pass

class DeterministicSIRsystem (SIRsystem):

    """Define a specialized subclass of the general SIRsystem for modeling SIR dynamics with a deterministic differential
    equation model"""

    def dydt(self, y, t):
        """Define the right-hand-side of the equation dy/dt = dydt(y,t) for use with odeint integrator;
        note that because this is defined as a method on the class rather than as a free-standing function, the first
        argument of the function must be the self instance rather than the current state vector y"""
        pass

    def run(self, T, dt=None):
        """Integrate the ODE for the deterministic model from time 0 to time T, starting with the initial values stored
        in the S,I,R state variables; story the result in self.trajectory"""
        pass

class StochasticSIRsystem (SIRsystem):

    """Define a specialized subclass of the general SIRsystem for modeling SIR dynamics as a stochastic, continuous
    time process, using the Gillespie method of continuous time Monte Carlo"""

    def step(self):
        """Implement one step of Gillespie's Direct Method based on current reaction rates: identify a reaction to fire
        next, as well as a time at which it will fire, and execute that reaction (similar to as was described in the
        StochasticCells exercise)."""
        pass

    def run(self, T=None, make_traj=True):
        """Run the Gillespie algorithm for stochastic simulation from time 0 to at least time T, starting with the initial
        values stored in the S,I,R state variables; story the result in self.trajectory if make_traj argument is
        set to True"""
        pass


def SimulateDeterministicOutbreakSize(N, R0_range=np.arange(0.,5.,0.1)):
    """For a given population size N and an array of basic reproductive numbers R0, integrate a
    DeterministicSIRsystem model for long enough time to allow the outbreak to die out, and then record the
    final size of the outbreak as a function of R0"""
    pass

def CalculateDeterministicOutbreakSize(N, R0_range=np.arange(0.,5.,0.1)):
    """For a given population size N and an array of basic reproductive numbers R0, solve an implicit equation
    for the final outbreak size using the fsolve root-finding routine.  Compare the simulated results found in
    SimulateDeterministicOutbreakSize with those solved for here"""
    pass

def StochasticOutbreakSizeDistribution(N, beta, gamma, Nruns):
    """For a given population size N and model parameters beta and gamma, execute Nruns dynamical runs of the
    StochasticSIRsystem to calculate the overall size of each outbreak (total number of hosts infected) after it has
    finally died out; store the outbreak sizes in an array, which is returned for further characterization."""
    pass

def FractionLargeOutbreaks(osd, N):
    """For a given array of outbreak sizes, with total population size N, calculate the fraction of those outbreaks
    that are 'large'.  Strictly speaking, we are intended in outbreaks whose sizes constitutes a finite fraction of the
    population in the limit of infinite population size, but for a finite system size N, we can implement a heuristic
    cutoff separate 'large' from 'small' outbreaks; determining such a cutoff can be assisted by examining the
    outbreak size distribution."""
    pass



# Copyright (C) Cornell University
# All rights reserved.
# Apache License, Version 2.0