network.py

from itertools import combinations

import networkx as nx

from .packet import PacketGenerator
from .node import NodeFactory
from .routers import types as routers
from .traces import types as traces

class Network:
    ''''''
    def __init__(self, env,
                 packets=None,
                 node_factory=None,
                 community=None,
                 trace=None):
        ''''''
        self.env = env

        # contact trace
        if trace is None:
            trace = traces['random']()
        self.trace = trace
        self.trace.start(env, network=self)

        # community detection
        self.community = community
        if community is not None:
            self.community.start(env, network=self)

        # packet generation
        if packets is None:
            packets = {}
        self.packets = PacketGenerator(**packets)
        self.packets.start(env, network=self)

        # create node network
        if node_factory is None:
            node_factory = NodeFactory(tick_rate=1, router=routers['direct'])
        self.nodes = [
            node_factory(self, nid)
            for nid in range(self.trace.nodes)
        ]
        self.links = [
            (a, b)
            for a, b in combinations(self.nodes, 2)
        ]

        # set up networkx graph
        self.graph = nx.Graph()
        self.graph.add_nodes_from(self.nodes)
        self.graph.add_edges_from([
            (a, b, { 'state': False })
            for a, b in self.links
        ])

    def set_link(self, a, b, state):
        if isinstance(a, int):
            a = self.nodes[a]

        if isinstance(b, int):
            b = self.nodes[b]

        edge = self[a][b]
        if edge['state'] == state:
            return

        if state is None:
            state = not edge['state']
        edge['state'] = state

        if self.community:
            self.community.set_link(a, b, state, self.env.now)

        if state:
            a.route_packets()
            b.route_packets()

    def toggle_link(self, a, b):
        self.set_link(a, b, None)

    def send_link(self, a, b, packet):
        ''''''
        if self.graph[a][b]['state']:
            # TODO: transfer delay
            b.recv(packet)
        else:
            raise Exception('Nodes {} and {} not connected'.format(a, b))

    def __getitem__(self, node):
        ''''''
        return self.graph[node]