<reponame>MattSegal/AuTuMN
from autumn.constants import Flow, Compartment


def add_infection_flows(working_flows, n_exposed):
    """
    Add standard infection flows for transition from susceptible to exposed through infection
    """
    working_flows += [{
        'type': Flow.INFECTION_FREQUENCY,
        'parameter': 'contact_rate',
        'origin': Compartment.SUSCEPTIBLE,
        'to': Compartment.EXPOSED + '_1' if n_exposed > 1 else Compartment.EXPOSED
    }]
    return working_flows


def add_transition_flows(list_of_flows, n_origin, n_to, origin_compartment, to_compartment, parameter_name):
    """
    Add flow from end of sequential exposed compartments to start of presymptomatic compartments
    """
    list_of_flows += [{
        'type': Flow.STANDARD,
        'origin': origin_compartment + '_' + str(n_origin) if n_origin > 1 else origin_compartment,
        'to': to_compartment + '_1' if n_to > 1 else to_compartment,
        'parameter': parameter_name
    }]
    return list_of_flows


def add_recovery_flows(working_flows, n_late):
    """
    Add standard recovery flows
    Differs from the transition flows in that the recovered compartment is never duplicated
    """
    working_flows += [{
        'type': Flow.STANDARD,
        'parameter': 'within_late',
        'to': Compartment.RECOVERED,
        'origin': Compartment.LATE_INFECTIOUS + '_' + str(n_late) if
        n_late > 1 else Compartment.LATE_INFECTIOUS
    }]
    return working_flows


def add_sequential_compartment_flows(working_flows, n_compartments, compartment_name):
    """
    Add standard flows for progression through any sequential compartments
    """
    for i_flow in range(1, n_compartments):
        working_flows += [{
            'type': Flow.STANDARD,
            'parameter': 'within_' + compartment_name,
            'origin': compartment_name + '_' + str(i_flow),
            'to': compartment_name + '_' + str(i_flow + 1)
        }]
    return working_flows


def add_infection_death_flows(working_flows, n_infectious):
    """
    Add infection-related death flows to the infectious compartments
    """
    for i_comp in range(n_infectious):
        working_flows.append({
            'type': Flow.COMPARTMENT_DEATH,
            'parameter': 'infect_death',
            'origin': Compartment.LATE_INFECTIOUS + '_' + str(i_comp + 1) if
            n_infectious > 1 else
            Compartment.LATE_INFECTIOUS
        })
    return working_flows

from autumn.constants import Flow, Compartment


def add_infection_flows(working_flows, n_exposed):
    """
    Add standard infection flows for transition from susceptible to exposed through infection
    """
    working_flows += [{
        'type': Flow.INFECTION_FREQUENCY,
        'parameter': 'contact_rate',
        'origin': Compartment.SUSCEPTIBLE,
        'to': Compartment.EXPOSED + '_1' if n_exposed > 1 else Compartment.EXPOSED
    }]
    return working_flows


def add_transition_flows(list_of_flows, n_origin, n_to, origin_compartment, to_compartment, parameter_name):
    """
    Add flow from end of sequential exposed compartments to start of presymptomatic compartments
    """
    list_of_flows += [{
        'type': Flow.STANDARD,
        'origin': origin_compartment + '_' + str(n_origin) if n_origin > 1 else origin_compartment,
        'to': to_compartment + '_1' if n_to > 1 else to_compartment,
        'parameter': parameter_name
    }]
    return list_of_flows


def add_recovery_flows(working_flows, n_late):
    """
    Add standard recovery flows
    Differs from the transition flows in that the recovered compartment is never duplicated
    """
    working_flows += [{
        'type': Flow.STANDARD,
        'parameter': 'within_late',
        'to': Compartment.RECOVERED,
        'origin': Compartment.LATE_INFECTIOUS + '_' + str(n_late) if
        n_late > 1 else Compartment.LATE_INFECTIOUS
    }]
    return working_flows


def add_sequential_compartment_flows(working_flows, n_compartments, compartment_name):
    """
    Add standard flows for progression through any sequential compartments
    """
    for i_flow in range(1, n_compartments):
        working_flows += [{
            'type': Flow.STANDARD,
            'parameter': 'within_' + compartment_name,
            'origin': compartment_name + '_' + str(i_flow),
            'to': compartment_name + '_' + str(i_flow + 1)
        }]
    return working_flows


def add_infection_death_flows(working_flows, n_infectious):
    """
    Add infection-related death flows to the infectious compartments
    """
    for i_comp in range(n_infectious):
        working_flows.append({
            'type': Flow.COMPARTMENT_DEATH,
            'parameter': 'infect_death',
            'origin': Compartment.LATE_INFECTIOUS + '_' + str(i_comp + 1) if
            n_infectious > 1 else
            Compartment.LATE_INFECTIOUS
        })
    return working_flows
Add standard infection flows for transition from susceptible to exposed through infection Add flow from end of sequential exposed compartments to start of presymptomatic compartments Add standard recovery flows     Differs from the transition flows in that the recovered compartment is never duplicated Add standard flows for progression through any sequential compartments Add infection-related death flows to the infectious compartments