import numpy as np
import networkx as nx
from math import gamma
from scipy.optimize import root_scalar

from mesa import Model
from mesa.time import RandomActivation, SimultaneousActivation
from mesa.datacollection import DataCollector

from scseirx.testing_strategy import Testing

## data collection functions ##
def get_N_diagnostic_tests(model):
    return model.number_of_diagnostic_tests


def get_N_preventive_screening_tests(model):
    return model.number_of_preventive_screening_tests


def get_infection_state(agent):
    if agent.exposed == True: return 'exposed'
    elif agent.infectious == True: return 'infectious'
    elif agent.recovered == True: return 'recovered'
    else: return 'susceptible'

def get_quarantine_state(agent):
    if agent.quarantined == True: return True
    else: return False


def get_undetected_infections(model):
    return model.undetected_infections


def get_predetected_infections(model):
    return model.predetected_infections


def get_pending_test_infections(model):
    return model.pending_test_infections


def get_diagnostic_test_detected_infections_student(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['student']
def get_diagnostic_test_detected_infections_teacher(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['teacher']
def get_diagnostic_test_detected_infections_family_member(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['family_member']
def get_diagnostic_test_detected_infections_resident(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['resident']
def get_diagnostic_test_detected_infections_employee(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['employee']
def get_diagnostic_test_detected_infections_unistudent(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['unistudent']
def get_diagnostic_test_detected_infections_lecturer(model):
    return model.positive_tests[model.Testing.diagnostic_test_type]['lecturer']

diagnostic_test_detected_infections_funcs = {
    'student':get_diagnostic_test_detected_infections_student,
    'teacher':get_diagnostic_test_detected_infections_teacher,
    'family_member':get_diagnostic_test_detected_infections_family_member,
    'resident':get_diagnostic_test_detected_infections_resident,
    'employee':get_diagnostic_test_detected_infections_employee,
    'unistudent':get_diagnostic_test_detected_infections_unistudent,
    'lecturer':get_diagnostic_test_detected_infections_lecturer
}

def get_preventive_test_detected_infections_student(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['student']
def get_preventive_test_detected_infections_teacher(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['teacher']
def get_preventive_test_detected_infections_family_member(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['family_member']
def get_preventive_test_detected_infections_resident(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['resident']
def get_preventive_test_detected_infections_employee(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['employee']
def get_preventive_test_detected_infections_unistudent(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['unistudent']
def get_preventive_test_detected_infections_lecturer(model):
    return model.positive_tests[model.Testing.preventive_screening_test_type]['lecturer']

preventive_test_detected_infections_funcs = {
    'student':get_preventive_test_detected_infections_student,
    'teacher':get_preventive_test_detected_infections_teacher,
    'family_member':get_preventive_test_detected_infections_family_member,
    'resident':get_preventive_test_detected_infections_resident,
    'employee':get_preventive_test_detected_infections_employee,
    'unistudent':get_preventive_test_detected_infections_unistudent,
    'lecturer':get_preventive_test_detected_infections_lecturer
}


# parameter sanity check functions


def check_positive(var):
	assert var >= 0, 'negative number'
	return var


def check_bool(var):
	assert type(var) == bool, 'not a bool'
	return var


def check_positive_int(var):
    if var == None:
        return var
    assert type(var) == int, 'not an integer'
    assert var >= 0, 'negative number'
    return var


def check_contact_type_dict(var):
	assert type(var) == dict, 'not a dictionary'
	assert set(var.keys()).issubset({'very_far', 'far', 'intermediate', 'close'}), \
		'does not contain the correct contact types (has to be very_far, far, intermediate or close)'
	assert all((isinstance(i, int) or isinstance(i, float)) for i in var.values()), \
		'contact type weights are not numeric'

	return var


def check_K1_contact_types(var):
    for area in var:
        assert area in ['very_far', 'far', 'intermediate',
            'close'], 'K1 contact type not recognised'
    return var


def check_testing(var):
    assert var in ['diagnostic', 'background', 'preventive',
                   'background+preventive', False], \
        'unknown testing mode: {}'.format(var)

    return var



def check_probability(var):
	assert (type(var) == float) or (var == 0) or (var == 1), \
		 '{} not a float'.format(var)
	assert var >= 0, 'probability negative'
	assert var <= 1, 'probability larger than 1'
	return var


def check_graph(var):
    assert type(var) in [nx.Graph, nx.MultiGraph], 'not a networkx graph'
    assert len(var.nodes) > 0, 'graph has no nodes'
    assert len(var.edges) > 0, 'graph has no edges'
    areas = [e[2]['contact_type'] for e in var.edges(data=True)]
    areas = set(areas)
    for a in areas:
        assert a in {'very_far', 'far', 'intermediate',
            'close'}, 'contact type {} not recognised'.format(a)
    return var


def check_index_case(var, agent_types):
	allowed_strings = agent_types[:]
	allowed_strings.extend(['continuous'])
	assert var in allowed_strings, 'unknown index case mode'
	return var


def check_discount(var):
    if var['slope'] != None:
        assert var['slope'] <= 0, 'slope needs to be <= 0 or None'
        assert np.abs(var['slope']) <= 1, 'absolute value of slope needs to be <= 1'
    assert var['intercept'], 'intercept needs to be positive'
    assert var['intercept'], 'intercept needs to be <= 1'
    return var


def get_weibull_shape(k, mu, var):
    '''
    Calculates the shape parameter of a Weibull distribution, given its mean
    mu and its variance var
    '''
    return var / mu**2 - gamma(1 + 2/k) / gamma(1+1/k)**2 + 1



def get_weibull_scale(mu, k):
    '''
    Calculates the scale parameter of a Weibull distribution, given its mean
    mu and its shape parameter k
    '''
    return mu / gamma(1 + 1/k)


def weibull_two_param(shape, scale):
    '''
    A two-parameter Weibull distribution, based on numpy ramdon's single
    parameter distribution. We use this distribution in the simulation to draw
    random epidemiological parameters for agents from the given distribution
    See https://numpy.org/doc/stable/reference/random/generated/numpy.random.weibull.html
    '''
    return scale * np.random.weibull(shape)


class SEIRX(Model):
    '''
    A model with a number of different agents that reproduces
    the SEIRX dynamics of pandemic spread in a facility. Note:
    all times are set to correspond to days

    G : networkx undirected graph
        Interaction graph between agents. Edges have to have edge the edge 
        attribute 'contact_type' specifying the closeness of contacts, which can
        be ['very far', 'far', 'intermediate' and 'close']. Nodes have to have 
        the node attribute 'type' which specifies the agent type of the given 
        node (for example 'student' or 'teacher' in a school scenario).
        In addition, nodes can have the attribute 'unit', which assigns them to 
        a unit in space (for example a 'class' in a school scenario).

    verbosity : integer in [0, 1, 2]
        Controls text output to std out to track simulation progress and 
        transmission dynamics. Default = 0.

    base_transmission_risk : float
        Base probability to transmit an infection durig a contact of type 
        "close" if no other measures or biological parameters influence the
        infection risk.

    testing : str
        Determines the testing strategy of the model. Default: 'diagnostic'.
        'diagnostic':   only diagnostic tests for symptomatic agents
        'background':   adds background screens of all agents after a positive
                        diagnostic test
        'preventive':   adds preventive screens of agent groups to diagnostic
                        testing. Screens happen in time intervals specified 
                        separately for each agent group in the variable 
                        'screening_interval'.
        'background+preventive': preventive screens AND background screens on
                        top of diagnostic testing.

    infection_duration : int or [float, float] 
        Parameter determining the infection duration. Default: 11 
        NOTE: includes the time an agent is exposed but not yet infectious at 
        the beginning of an infection.
        positive integer:   mean or median of the infection duration in days
        list of two floats: mean and standard deviation of a distribution
                            specifying the infection duration in days. These
                            numbers will be used to construct a Weibull
                            distribution from which the infection duration will
                            be drawn for every agent individually

    exposure_duration : int or [float, float] 
        Parameter determining the exposure duration. Default: 4. 
        Sets the time from transmission to becoming infectious
        positive integer:   mean or median of the exposure duration in days
        list of two floats: mean and standard deviation of a distribution
                            specifying the exposure duration in days. These
                            numbers will be used to construct a Weibull
                            distributoin from which the exposure duration will
                            be drawn for every agent individually.

    time_until_symptoms : int or [float, float] 
        Sets the time from transmission to (potentially) developing symptoms. 
        Default: 6. Symptom probability has to be set for each agent group 
        individually using the parameter 'symptom_probability'
        positive integer:   mean or median of the time until symptoms in days
        list of two floats: mean and standard deviation of a distribution
                            specifying the time until symptoms in days. These
                            numbers will be used to construct a Weibull
                            distribution from which the time until symptoms will
                            be drawn for every agent individually.

    quarantine_duration : int
        Positive integer, sets the time a positively tested agent is quarantined
        in days. Default: 14. 

    subclinical_modifier : float
        Modifies the infectiousness of asymptomatic cases. Example: if 
        subclinical_modifier = 0.5, the infectiousness of an asymptomatic case 
        will be reduced to 50%. Default: 1.0.

    infection_risk_contact_type_weights: dictionary 
        Has to have the form
        {
        'very_far':    float, 
        'far':         float, 
        'intermediate':float, 
        'close':       float
        }
    Sets transmission risk multipliers for different contact types of agents 
    specified in the contact network G. Default: {'very_far': 0.1,
    'far': 0.5, 'intermediate': 1, 'close': 3}


    K1_contact_types : list of strings
        Definition of contact types for which agents are considered "K1 contact 
        persons" if they had contact to a positively tested person with a 
        specified contact intensity. The list may only include strings from the
        list ['very_far', 'far', 'intermediate', 'close'], i.e. the keys of the
        infection_risk_contact_type_weights dictionary.  Default = ['close'].

    diagnostic_test_type : str 
        Specifies the test technology and test result turnover time used for 
        diagnostic testing. For example 'same_day_antigen' or 'two_day_PCR'. 
        See module "Testing" for different implemented testing techologies.
        Default: 'one_day_PCR'. 

    preventive_screening_test_type : str
        Specifies the test technology and test result turnover time used for 
        preventive sreening. For example 'same_day_antigen' or 'two_day_PCR'. 
        See module "Testing" for different implemented testing techologies.
        Default: 'one_day_PCR'.

    follow_up_testing_interval : int or None
        Sets the time a follow-up screen (background screen) is initiated after 
        an initial screen triggered by a positive test result. Only applies if 
        the testing strategy is 'background' or 'preventive'. If "None" is 
        chosen, no background screens are performed. Default: None. 

    liberating_testing : bool 
        Flag that specifies, whether or not an agent is released from quarantine 
        after returning a negative test result. Default: False.

	index_case: str
        Specifies how infections are introduced into the facility. Can bei 
        either a scenario-specific agent type (for example "student" or 
        "teacher" in the school scenario) or "continuous". 
        agent_type:     If an agent type is specified, a single randomly chosen 
                        agent from this agent group will become the index case 
                        and no further index cases will be introduced into the 
                        scenario.
        'continuous':   In this case, agents have a continuous risk to become
                        index cases in every simulation step. The risk has to
                        be specified for every agent group individually, using
                        the 'index_probability' parameter. If only a single
                        agent group has a non-zero index probability, then only
                        agents from this group can become index cases.
        Default = 'employee' (nursing home scenario), 'teacher' (school 
        scenario), 'lecturer' (university scenario). 


    agent_types: dictionary of dictionaries
        Specifies the parameters 'screening_interval', 'index_probability', 
        'mask' and 'vaccination_probability' individually for every agent group
        in the model. The dictionary therefore needs to have an entry for every
        agent type and every entry needs to have the following form:
        {
        agent_type:
            {
            screening_interval : int
                Number of days between each preventive screen in this agent 
                group.
            index_probability : float 
                Probability to become an index case in each time step. Needs to
                be in in the range [0, 1]. 
            mask : bool
                Whether or not the agent type is wearing a mask.
            vaccination_probability : float
                Probability of an agent in the agent group to be vaccinated.
                Needs to be in the range [0, 1].
            }
        }
    The dictionary's keys are the names of the agent types which have to
    correspond to the node attributes in the contact graph. 

    age_transmission_risk_discount : dictionary
        Specifies the linear relationship between agent age and transmission 
        risk. Applies both to the risk of transmitting and the risk of receiving 
        an infection. This is only used in scenarios in which agents have an age
        attribute (school). The dictionary needs to have the following fields:
            'slope' : float
                The slope of the linear relationship.
            'intercept' : float
                The intercept of the linear relationship.
        Default: {'slope':-0.02, 'intercept:1'}


    age_symptom_modification : dictionary
        Specifies the linear relationship between agent age and the probability
        to develop symptoms. This is only used in scenarios in which agents have
        an age attribute (school). The dictionary needs to have the following
        fields:
            'slope' : float
                The slope of the linear relationship.
            'intercept' : float
                The intercept of the linear relationship.
        Default: {'slope':-0.02545, 'intercept:0.854545'}

    mask_filter_efficiency : dictionary
        Specifies the effectiveness of masks to prevent the transmission and
        reception of an infection. The dictionary needs to have the following
        fields:
            'inhale' : float
                Reduction of the risk to become infected if the receiving agent
                wears a mask.
            'exhale' : float
                Reduction of the risk to transmit an infection if the 
                transmitting agent wears a mask.
        Default: {'inhale':0, 'exhale':0}

    transmission_risk_ventilation_modifier : float
        Reduction of the transmission risk due to ventilation. Default: 0.

    transmission_risk_vaccination_modifier : dictionary
        Reduction of the transmission and reception risk due to vaccination.
        The dictionary needs to have the following fields:
            'reception' : float
                Reduction of the probability to get infected if the infectee is
                vaccinated.
            'transmission' : float
                Reduction of the probability to transmit the infection if the
                infected agent is vaccinated.
        Default: {'reception':1, 'transmission':0}

    N_days_in_network : int
        Number of distinct days the graph object contains. If the simulation
        duration (in days) surpasses the number of days in the graph, the
        simulation loops back to the first day etc. Default: 7.

    seed : int
        Fixes the seed of the simulation to enable repeatable simulation runs. 
        If seed = None, the simulation will be initialized at random.
        Default: None.
    '''

    def __init__(self, G,
        verbosity = 0,
        base_transmission_risk = 0.05,
        testing='diagnostic',
        exposure_duration = [5.0, 1.9],
        time_until_symptoms = [6.4, 0.8],
        infection_duration = [10.91, 3.95],
        quarantine_duration = 10,
        subclinical_modifier = 0.6,
        infection_risk_contact_type_weights = {
            'very_far': 0.1,
            'far': 0.25,
            'intermediate': 0.5,
            'close': 1},
        K1_contact_types = ['close'],
        diagnostic_test_type = 'one_day_PCR',
        preventive_screening_test_type = 'same_day_antigen',
        follow_up_testing_interval = None,
        liberating_testing = False,
        index_case = 'teacher',
        agent_types = {
            'teacher':      {'screening_interval': None,
                             'index_probability': 0,
                             'mask':False,
                             'vaccination_ratio': 0},
            'student':      {'screening_interval': None,
                             'index_probability': 0,
                             'mask':False,
                             'vaccination_ratio': 0},
            'family_member':{'screening_interval': None,
                             'index_probability': 0,
                             'mask':False,
                             'vaccination_ratio': 0}},
        age_transmission_risk_discount = \
             {'slope':-0.02,
              'intercept':1},
        age_symptom_modification = \
             {'slope':-0.02545,
              'intercept':0.854545},
        mask_filter_efficiency = {'exhale':0, 'inhale':0},
        transmission_risk_ventilation_modifier = 0,
        transmission_risk_vaccination_modifier = {
            'reception':1,
            'transmission':0},
        N_days_in_network = 7,
        seed = None):

        # mesa models already implement fixed seeds through their own random
        # number generations. Sadly, we need to use the Weibull distribution
        # here, which is not implemented in mesa's random number generation
        # module. Therefore, we need to initialize the numpy random number
        # generator with the given seed as well
        if seed != None:
            np.random.seed(seed)

        # sets the (daily) transmission risk for a household contact without
        # any precautions. Target infection ratios are taken from literature
        # and the value of the base_transmission_risk is calibrated such that
        # the simulation produces the correct infection ratios in a household
        # setting with the given distributions for epidemiological parameters
        # of agents
        self.base_transmission_risk = base_transmission_risk
    	# sets the level of detail of text output to stdout (0 = no output)
        self.verbosity = check_positive_int(verbosity)
        # flag to turn off the testing & tracing strategy
        self.testing = check_testing(testing)
        self.running = True  # needed for the batch runner implemented by mesa
        # set the interaction mode to simultaneous activation
        self.schedule = SimultaneousActivation(self)


        # internal step counter used to launch screening tests
        self.Nstep = 0

        # since we may have day-specific contact networks, we need
        # to keep track of the day of the week. Since the index case
        # per default is introduced at step 0 in index case mode, we
        # need to offset the starting day by a random number of days
        # to prevent artifacts from always starting on the same day of the week
        self.N_days_in_network = N_days_in_network
        self.day_offset = self.random.randint(1, 8)
        self.day = self.Nstep + self.day_offset

        ## epidemiological parameters: can be either a single integer or the
        # mean and standard deviation of a distribution
        self.epi_params = {}
        # counter to track the number of pathological parameter combinations
        # that had to be re-rolled (only here for debugging and control reasons)
        self.param_rerolls = 0

        for param, param_name in zip([exposure_duration, time_until_symptoms,
                infection_duration],['exposure_duration', 'time_until_symptoms',
                'infection_duration']):

            if isinstance(param, int):
                self.epi_params[param_name] = check_positive_int(param)

            elif isinstance(param, list) and len(param) == 2:

                mu = check_positive(param[0])
                var = check_positive(param[1]**2)
                shape = root_scalar(get_weibull_shape, args=(mu, var),
                            method='toms748', bracket=[0.2, 500]).root
                scale = get_weibull_scale(mu, shape)

                self.epi_params[param_name] = [shape, scale]
            else:
                print('{} format not recognized, should be either a single '+\
                  'int or a tuple of two positive numbers'.format(param_name))


        # duration of quarantine
        self.quarantine_duration = check_positive_int(quarantine_duration)

        self.infection_risk_area_weights = check_contact_type_dict(
            infection_risk_contact_type_weights)

        # modifier for infectiosness for asymptomatic cases
        self.subclinical_modifier = check_positive(subclinical_modifier)
        # modifiers for the infection risk, depending on contact type
        self.infection_risk_contact_type_weights = infection_risk_contact_type_weights

        # modifications for age-dependent transmission and reception risks and
        # symptom probabilities
        self.age_transmission_risk_discount = \
            check_discount(age_transmission_risk_discount)

        self.age_symptom_modification = age_symptom_modification
            #check_discount(age_symptom_modification)

        self.mask_filter_efficiency = mask_filter_efficiency
        self.transmission_risk_ventilation_modifier = \
            transmission_risk_ventilation_modifier
        self.transmission_risk_vaccination_modifier = \
            transmission_risk_vaccination_modifier
        ## agents and their interactions
        # interaction graph of agents
        self.G = check_graph(G)
        # add weights as edge attributes so they can be visualised easily
        if type(self.G) == nx.MultiGraph:
            for (u, v, key, contact_type) in self.G.edges(keys=True,
                    data='contact_type'):
                self.G[u][v][key]['weight'] = \
                    self.infection_risk_contact_type_weights[contact_type]
        else:
            for e in G.edges(data=True):
                G[e[0]][e[1]]['weight'] = self.infection_risk_contact_type_weights\
                	[G[e[0]][e[1]]['contact_type']]

        # extract the different agent types from the contact graph
        self.agent_types = list(agent_types.keys())
        # dictionary of available agent classes with agent types and classes
        self.agent_classes = {}
        if 'resident' in agent_types:
            from scseirx.agent_resident import resident
            self.agent_classes['resident'] = resident
        if 'employee' in agent_types:
            from scseirx.agent_employee import employee
            self.agent_classes['employee'] = employee
        if 'student' in agent_types:
            from scseirx.agent_student import student
            self.agent_classes['student'] = student
        if 'teacher' in agent_types:
            from scseirx.agent_teacher import teacher
            self.agent_classes['teacher'] = teacher
        if 'family_member' in agent_types:
            from scseirx.agent_family_member import family_member
            self.agent_classes['family_member'] = family_member
        if 'lecturer' in agent_types:
            from scseirx.agent_lecturer import lecturer
            self.agent_classes['lecturer'] = lecturer
        if 'unistudent' in agent_types:
            from scseirx.agent_unistudent import unistudent
            self.agent_classes['unistudent'] = unistudent

        ## set agent characteristics for all agent groups
        # list of agent characteristics
        params = ['screening_interval','index_probability', 'mask' ,'vaccination_ratio',
                  'voluntary_testing_rate']

        # default values that are used in case a characteristic is not specified
        # for an agent group
        defaults = {'screening_interval':None,
                    'index_probability':0,
                    'mask':False,
                    'vaccination_ratio':0,
                    'voluntary_testing_rate':1
                    }

        # sanity checks that are applied to parameters passed to the class
        # constructor to make sure they conform to model expectations
        check_funcs = [check_positive_int, check_probability, check_bool,
                       check_probability, check_probability]

        # member dicts that store the parameter values for each agent group
        self.screening_intervals = {}
        self.index_probabilities = {}
        self.masks = {}
        self.vaccination_probabilities = {}
        self.voluntary_testing_rates = {}


        param_dicts = [self.screening_intervals, self.index_probabilities,
                    self.masks, self.vaccination_probabilities, self.voluntary_testing_rates]

        # iterate over all possible agent parameters and agent groups: set the
        # respective value to the value passed through the constructor or to
        # the default value if no value has been passed
        for param,param_dict,check_func in zip(params,param_dicts,check_funcs):
            for at in self.agent_types:
                try:
                    param_dict.update({at:check_func(agent_types[at][param])})
                except KeyError:
                    param_dict.update({at:defaults[param]})

        # pass all parameters relevant for the testing strategy to the testing
        # class. NOTE: this separation is not a strictly necessary design
        # decision but I like to keep the parameters related to testing and
        # tracing in a separate place
        self.Testing = Testing(self, diagnostic_test_type,
             preventive_screening_test_type,
             check_positive_int(follow_up_testing_interval),
             self.screening_intervals,
             check_bool(liberating_testing),
             check_K1_contact_types(K1_contact_types),
             verbosity)


        # specifies either continuous probability for index cases in agent
        # groups based on the 'index_probability' for each agent group, or a
        # single (randomly chosen) index case in the passed agent group
        self.index_case = check_index_case(index_case, self.agent_types)

        self.num_agents = {}

        ## add agents
        # extract the agent nodes from the graph and add them to the scheduler
        for agent_type in self.agent_types:
            IDs = [x for x,y in G.nodes(data=True) if y['type'] == agent_type]
            self.num_agents.update({agent_type:len(IDs)})

            # get the agent locations (units) from the graph node attributes
            units = [self.G.nodes[ID]['unit'] for ID in IDs]

            # determine the agents that will be vaccinated, given the 
            # vaccination ratio of the respective agent group
            vaccination_status = np.asarray([False] * len(IDs))
            if self.vaccination_probabilities[agent_type] > 0:
                n = round(self.vaccination_probabilities[agent_type] * len(IDs))
                idx = list(range(len(IDs)))
                rnd_idx = np.asarray(self.random.sample(idx, n))
                vaccination_status[rnd_idx] = True


            for ID, unit, vaccinated in zip(IDs, units, vaccination_status):

                tmp_epi_params = {}
                # for each of the three epidemiological parameters, check if
                # the parameter is an integer (if yes, pass it directly to the
                # agent constructor), or if it is specified by the shape and
                # scale parameters of a Weibull distribution. In the latter
                # case, draw a new number for every agent from the distribution
                # NOTE: parameters drawn from the distribution are rounded to
                # the nearest integer
                while True:
                    for param_name, param in self.epi_params.items():
                        if isinstance(param, int):
                            tmp_epi_params[param_name] = param

                        else:
                            tmp_epi_params[param_name] = \
                                round(weibull_two_param(param[0], param[1]))

                    if tmp_epi_params['exposure_duration'] > 0 and \
                       tmp_epi_params['time_until_symptoms'] >= \
                       tmp_epi_params['exposure_duration'] and\
                       tmp_epi_params['infection_duration'] > \
                       tmp_epi_params['exposure_duration']:
                       break
                    else:
                        self.param_rerolls += 1
                        if verbosity > 1:
                            print('pathological epi-param case found!')
                            print(tmp_epi_params)

                # check if the agent participates in voluntary testing
                p = self.voluntary_testing_rates[agent_type]
                voluntary_testing = np.random.choice([True, False],
                         p=[p, 1-p])

                # construct the agent object
                a = self.agent_classes[agent_type](ID, unit, self,
                    tmp_epi_params['exposure_duration'],
                    tmp_epi_params['time_until_symptoms'],
                    tmp_epi_params['infection_duration'],
                    vaccinated,
                    voluntary_testing,
                    verbosity)
                self.schedule.add(a)


		# infect the first agent in single index case mode
        if self.index_case != 'continuous':
            infection_targets = [
                a for a in self.schedule.agents if a.type == index_case]
            # pick a random agent to infect in the selected agent group
            target = self.random.randint(0, len(infection_targets) - 1)
            infection_targets[target].exposed = True
            if self.verbosity > 0:
                print('{} exposed: {}'.format(index_case,
                    infection_targets[target].ID))


        # list of agents that were tested positive this turn
        self.newly_positive_agents = []
        # flag that indicates if there were new positive tests this turn
        self.new_positive_tests = False
        # dictionary of flags that indicate whether a given agent group has
        # been creened this turn
        self.screened_agents= {
            'reactive':{agent_type: False for agent_type in self.agent_types},
            'follow_up':{agent_type: False for agent_type in self.agent_types},
            'preventive':{agent_type: False for agent_type in self.agent_types}}


        # dictionary of counters that count the days since a given agent group
        # was screened. Initialized differently for different index case modes
        if (self.index_case == 'continuous') or \
      	   (not np.any(list(self.Testing.screening_intervals.values()))):
        	self.days_since_last_agent_screen = {agent_type: 0 for agent_type in
        	self.agent_types}
        # NOTE: if we initialize these variables with 0 in the case of a single
        # index case, we introduce a bias since in 'single index case mode' the
        # first index case will always become exposed in step 0. To realize
        # random states of the preventive sceening procedure with respect to the
        # incidence of the index case, we have to randomly pick the days since
        # the last screen for the agent group from which the index case is
        else:
        	self.days_since_last_agent_screen = {}
        	for agent_type in self.agent_types:
        		if self.Testing.screening_intervals[agent_type] != None:
        			self.days_since_last_agent_screen.update({
        				agent_type: self.random.choice(range(0,
        				 self.Testing.screening_intervals[agent_type] + 1))})
        		else:
        			self.days_since_last_agent_screen.update({agent_type: 0})

        # dictionary of flags that indicates whether a follow-up screen for a
        # given agent group is scheduled
        self.scheduled_follow_up_screen = {agent_type: False for agent_type in
        	self.agent_types}

        # counters
        self.number_of_diagnostic_tests = 0
        self.number_of_preventive_screening_tests = 0
        self.positive_tests = {self.Testing.preventive_screening_test_type:
                            {agent_type:0 for agent_type in self.agent_types},
                               self.Testing.diagnostic_test_type:
                            {agent_type:0 for agent_type in self.agent_types}}

        self.undetected_infections = 0
        self.predetected_infections = 0
        self.pending_test_infections = 0
        self.quarantine_counters = {agent_type:0 for agent_type in agent_types.keys()}
        self.false_negative = 0

        # data collectors to save population counts and agent states every
        # time step

        model_reporters = {
                'N_diagnostic_tests':get_N_diagnostic_tests,
                'N_preventive_screening_tests':get_N_preventive_screening_tests,
                'undetected_infections':get_undetected_infections,
                'predetected_infections':get_predetected_infections,
                'pending_test_infections':get_pending_test_infections
                }

        for agent_type in self.agent_types:
            model_reporters.update({
                'diagnostic_test_detected_infections_{}'.format(agent_type):\
                diagnostic_test_detected_infections_funcs[agent_type]
                })
            model_reporters.update({
                'preventive_test_detected_infections_{}'.format(agent_type):\
                preventive_test_detected_infections_funcs[agent_type]
                })


        self.datacollector = DataCollector(
            model_reporters=model_reporters,
            agent_reporters=
            	{
            	'infection_state': get_infection_state,
                'quarantine_state': get_quarantine_state
                })


    ## transmission risk modifiers
    def get_transmission_risk_contact_type_modifier(self, source, target):
        # construct the edge key as combination between agent IDs and day
        n1 = source.ID
        n2 = target.ID
        tmp = [n1, n2]
        tmp.sort()
        n1, n2 = tmp
        key = '{}{}d{}'.format(n1, n2, self.day)
        contact_weight = self.G.get_edge_data(n1, n2, key)['weight']

        # the link weight is a multiplicative modifier of the link strength.
        # contacts of type "close" have, by definition, a weight of 1. Contacts
        # of type intermediate, far or very far have a weight < 1 and therefore
        # are less likely to transmit an infection. For example, if the contact
        # type far has a weight of 0.2, a contact of type far has only a 20%
        # chance of transmitting an infection, when compared to a contact of
        # type close. To calculate the probability of success p in the Bernoulli
        # trial, we need to reduce the base risk (or base probability of success)
        # by the modifications introduced by preventive measures. These
        # modifications are formulated in terms of "probability of failure", or
        # "q". A low contact weight has a high probability of failure, therefore
        # we return q = 1 - contact_weight here.
        q1 = 1 - contact_weight

        return q1


    def get_transmission_risk_age_modifier_transmission(self, source):
        '''linear function such that at age 18 the risk is that of an adult (=1).
        The slope of the line needs to be calibrated.
        '''
        age = source.age
        max_age = 18
        if age <= max_age:
            age_weight = self.age_transmission_risk_discount['slope'] * \
                 np.abs(age - max_age) + self.age_transmission_risk_discount['intercept']

            # The age weight can be interpreted as multiplicative factor that
            # reduces the chance for transmission with decreasing age. The slope
            # of the age_transmission_discount function is the decrease (in % of
            # the transmission risk for an 18 year old or above) of transmission
            # risk with every year a person is younger than 18 (the intercept is
            # 1 by definition).
            # To calculate the probability of success p in the Bernoulli
            # trial, we need to reduce the base risk (or base probability of 
            # success) by the modifications introduced by preventive measures. 
            # These modifications are formulated in terms of "probability of 
            # failure", or "q". A low age weight has a high probability of 
            # failure, therefore we return q = 1 - age_weight here.
            q2 = 1 - age_weight
        else:
            q2 = 0

        return q2


    def get_transmission_risk_age_modifier_reception(self, target):
        '''linear function such that at age 18 the risk is that of an adult (=1).
        The slope of the line needs to be calibrated.
        '''
        age = target.age
        max_age = 18
        if age <= max_age:
            age_weight = self.age_transmission_risk_discount['slope'] * \
            np.abs(age - max_age) + self.age_transmission_risk_discount['intercept']
            # see description in get_transmission_risk_age_modifier_transmission
            q3 = 1 - age_weight
        else:
            q3 = 0

        return q3


    # infectiousness is constant and high until symptom onset and then
    # decreases monotonically until agents are not infectious anymore
    # at the end of the infection_duration
    def get_transmission_risk_progression_modifier(self, source):
        if source.days_since_exposure < source.exposure_duration:
            progression_weight = 0
        elif source.days_since_exposure <= source.time_until_symptoms:
            progression_weight = 1
        elif source.days_since_exposure > source.time_until_symptoms and \
             source.days_since_exposure <= source.infection_duration:
            # we add 1 in the denominator, such that the source is also
            # (slightly) infectious on the last day of the infection_duration
            progression_weight = \
                 (source.days_since_exposure - source.time_until_symptoms) / \
                 (source.infection_duration - source.time_until_symptoms + 1)
        else:
            progression_weight = 0
        # see description in get_transmission_risk_age_modifier_transmission
        q4 = 1 - progression_weight

        return q4

    def get_transmission_risk_subclinical_modifier(self, source):
        if source.symptomatic_course == False:
            subclinical_weight = self.subclinical_modifier
        else:
            subclinical_weight = 1
        # see description in get_transmission_risk_age_modifier_transmission
        q5 = 1 - subclinical_weight
        return q5

    def get_transmission_risk_exhale_modifier(self, source):
        if source.mask:
            exhale_weight = self.mask_filter_efficiency['exhale']
        else:
            exhale_weight = 1
        # see description in get_transmission_risk_age_modifier_transmission
        q6 = 1 - exhale_weight
        return q6


    def get_transmission_risk_inhale_modifier(self, target):
        if target.mask:
            inhale_weight = self.mask_filter_efficiency['inhale']
        else:
            inhale_weight = 1
        # see description in get_transmission_risk_age_modifier_transmission
        q7 = 1 - inhale_weight
        return q7


    def get_transmission_risk_ventilation_modifier(self):
        ventilation_weight = self.transmission_risk_ventilation_modifier
        # see description in get_transmission_risk_age_modifier_transmission
        q8 = 1 - ventilation_weight
        return q8

    def get_transmission_risk_vaccination_modifier_reception(self, a):
        if a.vaccinated:
            q9 = self.transmission_risk_vaccination_modifier['reception']
        else:
            q9 = 0
        return q9

    def get_transmission_risk_vaccination_modifier_transmission(self, a):
        if a.vaccinated:
            q10 = self.transmission_risk_vaccination_modifier['transmission']
        else:
            q10 = 0
        return q10

    def test_agent(self, a, test_type):
        a.tested = True
        a.pending_test = test_type
        if test_type == self.Testing.diagnostic_test_type:
            self.number_of_diagnostic_tests += 1
        else:
            self.number_of_preventive_screening_tests += 1

        if a.exposed:
            # tests that happen in the period of time in which the agent is
            # exposed but not yet infectious. 
            # Note: tests[test_type]['time_until_testable'] is negative for
            # tests that can detect an infection before agents become infectious
            if a.days_since_exposure >= a.exposure_duration + \
                    self.Testing.tests[test_type]['time_until_testable']:
                
                if self.verbosity > 1:
                    print('{} {} sent positive sample (even though not infectious yet)'
                    .format(a.type, a.ID))
                a.sample = 'positive'
                self.predetected_infections += 1
                self.positive_tests[test_type][a.type] += 1
            else:
                if self.verbosity > 1: print('{} {} sent negative sample'
                    .format(a.type, a.ID))
                a.sample = 'negative'

        elif a.infectious:
            # tests that happen in the period of time in which the agent is
            # infectious and the infection is detectable by a given test
            # Note: tests[test_type]['time_until_testable'] is negative for 
            # tests that can detect an infection before agents become 
            # infectious. tests[test_type]['time_testable'] is negative for
            # tests that cease to detect an infection before agents stop being
            # infectious
            if a.days_since_exposure >= a.exposure_duration + \
                    self.Testing.tests[test_type]['time_until_testable'] and \
               a.days_since_exposure <= a.infection_duration + \
                    self.Testing.tests[test_type]['time_testable']:
                if self.verbosity > 1:
                    print('{} {} sent positive sample'.format(a.type, a.ID))
                a.sample = 'positive'
                self.positive_tests[test_type][a.type] += 1

            # track the undetected infections to assess how important they are
            # for infection spread
            else:
                if self.verbosity > 1:
                    print('{} {} sent negative sample (even though infectious)'
                    .format(a.type, a.ID))
                a.sample = 'negative'
                self.undetected_infections += 1

        else:
            if self.verbosity > 1: print('{} {} sent negative sample'
                .format(a.type, a.ID))
            a.sample = 'negative'

        # for same-day testing, immediately act on the results of the test
        if a.days_since_tested >= self.Testing.tests[test_type]['time_until_test_result']:
            a.act_on_test_result()

    def screen_agents(self, agent_group, test_type, screen_type):
        # only test agents that have not been tested already in this simulation
        # step and that are not already known positive cases

        if self.verbosity > 0:
            print('initiating {} {} screen'\
                                .format(screen_type, agent_group))

        untested_agents = [a for a in self.schedule.agents if
            (a.tested == False and a.known_positive == False
                and a.type == agent_group)]

        if len(untested_agents) > 0:
            self.screened_agents[screen_type][agent_group] = True
            self.days_since_last_agent_screen[agent_group] = 0

            # only test agents if they participate in voluntary testing
            if screen_type == 'preventive':
                for a in untested_agents:
                    if a.voluntary_testing:
                        self.test_agent(a, test_type)
                    else:
                        if self.verbosity > 1:
                            print('not testing {} {}, not participating in voluntary testing'\
                                .format(agent_group, a.ID))
            else:
                for a in untested_agents:
                    self.test_agent(a, test_type)

            if self.verbosity > 0:
                print()
        else:
            if self.verbosity > 0:
                print('no agents tested because all agents have already been tested')

    # the type of the test used in the pending test result is stored in the
    # variable pending_test

    def collect_test_results(self):
        agents_with_test_results = [a for a in self.schedule.agents if
            (a.pending_test and
             a.days_since_tested >= self.Testing.tests[a.pending_test]['time_until_test_result'])]

        return agents_with_test_results

    def trace_contacts(self, a):
        if a.quarantined == False:
            a.quarantined = True
            a.quarantine_start = self.Nstep

            if self.verbosity > 0:
                print('qurantined {} {}'.format(a.type, a.ID))

        # find all agents that share edges with the agent
        # that are classified as K1 contact types in the testing
        # strategy
        if a in self.G.nodes():
            K1_contacts = [e[1] for e in self.G.edges(a.ID, data=True) if
                e[2]['contact_type'] in self.Testing.K1_contact_types]
            K1_contacts = [a for a in self.schedule.agents if a.ID in K1_contacts]

            for K1_contact in K1_contacts:
                if self.verbosity > 0:
                    print('quarantined {} {} (K1 contact of {} {})'
                        .format(K1_contact.type, K1_contact.ID, a.type, a.ID))
                K1_contact.quarantined = True
                K1_contact.quarantine_start = self.Nstep

    def test_symptomatic_agents(self):
        # find symptomatic agents that have not been tested yet and are not
        # in quarantine and test them
        newly_symptomatic_agents = np.asarray([a for a in self.schedule.agents
            if (a.symptoms == True and a.tested == False and a.quarantined == False)])

        for a in newly_symptomatic_agents:
            # all symptomatic agents are quarantined by default
            if self.verbosity > 0:
                print('quarantined: {} {}'.format(a.type, a.ID))
            a.quarantined = True
            a.quarantine_start = self.Nstep

            self.test_agent(a, self.Testing.diagnostic_test_type)

    def quarantine_contacts(self):
        # trace and quarantine contacts of newly positive agents
        if len(self.newly_positive_agents) > 0:
            if self.verbosity > 0: print('new positive test(s) from {}'
                .format([a.ID for a in self.newly_positive_agents]))

            # send all K1 contacts of positive agents into quarantine
            for a in self.newly_positive_agents:
                self.trace_contacts(a)

            # indicate that a screen should happen because there are new
            # positive test results
            self.new_positive_tests = True
            self.newly_positive_agents = []

        else:
            self.new_positive_tests = False


    def step(self):
        self.day = (self.Nstep + self.day_offset) % self.N_days_in_network + 1
        # if the connection graph is time-resloved, set the graph that is
        # used to determine connections in this step to the sub-graph corres-
        # ponding to the current day of the week
        if self.dynamic_connections:
            self.G = self.day_connections[self.day]

        if self.verbosity > 0:
            print('day {}'.format(self.day))

        if self.testing:
            for agent_type in self.agent_types:
                for screen_type in ['reactive', 'follow_up', 'preventive']:
                    self.screened_agents[screen_type][agent_type] = False

            if self.verbosity > 0:
                print('* testing and tracing *')

            self.test_symptomatic_agents()


            # collect and act on new test results
            agents_with_test_results = self.collect_test_results()
            for a in agents_with_test_results:
                a.act_on_test_result()

            self.quarantine_contacts()

            # screening:
            # a screen should take place if
            # (a) there are new positive test results
            # (b) as a follow-up screen for a screen that was initiated because
            # of new positive cases
            # (c) if there is a preventive screening policy and it is time for
            # a preventive screen in a given agent group

            # (a)
            if (self.testing == 'background' or self.testing == 'background+preventive')\
               and self.new_positive_tests == True:
                for agent_type in self.screening_agents:
                    self.screen_agents(
                        agent_type, self.Testing.diagnostic_test_type, 'reactive')
                    self.scheduled_follow_up_screen[agent_type] = True

            # (b)
            elif (self.testing == 'background' or self.testing == 'background+preventive') and \
                self.Testing.follow_up_testing_interval != None and \
                sum(list(self.scheduled_follow_up_screen.values())) > 0:
                for agent_type in self.screening_agents:
                    if self.scheduled_follow_up_screen[agent_type] and\
                       self.days_since_last_agent_screen[agent_type] >=\
                       self.Testing.follow_up_testing_interval:
                        self.screen_agents(
                            agent_type, self.Testing.diagnostic_test_type, 'follow_up')
                    else:
                        if self.verbosity > 0:
                            print('not initiating {} follow-up screen (last screen too close)'\
                                .format(agent_type))

            # (c) 
            elif (self.testing == 'preventive' or self.testing == 'background+preventive')and \
                np.any(list(self.Testing.screening_intervals.values())):

                for agent_type in self.screening_agents:
                    interval = self.Testing.screening_intervals[agent_type]
                    assert interval in [7, 3, 2, None], \
                        'testing interval {} for agent type {} not supported!'\
                        .format(interval, agent_type)

                    # (c.1) testing every 7 days = testing on Mondays
                    if interval == 7 and self.day % interval == 1:
                        self.screen_agents(agent_type,
                            self.Testing.preventive_screening_test_type,\
                             'preventive')
                    # (c.2) testing every 3 days = testing on Mo & Turs
                    elif interval == 3 and self.day % interval in [1, 4]:
                            self.screen_agents(agent_type,
                            self.Testing.preventive_screening_test_type,\
                             'preventive')
                    # (c.3) testing every 2 days = testing on Mo, Wed & Fri
                    elif interval == 2 and self.day % interval in [1, 3, 5]:
                            self.screen_agents(agent_type,
                            self.Testing.preventive_screening_test_type,\
                             'preventive')
                    # No interval specified = no testing, even if testing
                    # mode == preventive
                    elif interval == None:
                        pass
                    else:
                        if self.verbosity > 0:
                            print('not initiating {} preventive screen (wrong day)'\
                                    .format(agent_type))
            else:
                # do nothing
                pass

            for agent_type in self.agent_types:
                if not (self.screened_agents['reactive'][agent_type] or \
                        self.screened_agents['follow_up'][agent_type] or \
                        self.screened_agents['preventive'][agent_type]):
                        self.days_since_last_agent_screen[agent_type] += 1


        if self.verbosity > 0: print('* agent interaction *')
        self.datacollector.collect(self)
        self.schedule.step()
        self.Nstep += 1
