network.py

from itertools import combinations

import networkx as nx

from .core import TickProcess

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]