from random import Random

from pydtn import Network, Node, random_trace, random_traffic

class RandomNode(Node):
    def __init__(self, seed=None, **options):
        super().__init__(**options)
        self.random = Random(seed)

    def forward(self, packet):
        neighbours = list(self.neighbours)
        if neighbours:
            target = self.random.choice(neighbours)
            return {target: 'random'}
        return {}

def main():
    nodes = 50
    seed  = 42
    traffic_speed = 1

    node_options = {
        'seed': seed,
        'tick_rate': 1,
    }

    nodes = {
        node_id: RandomNode(**node_options)
        for node_id in range(nodes)
    }

    trace = random_trace(nodes, seed=seed)
    traffic = random_traffic(nodes, speed=traffic_speed, seed=seed)
    network = Network(nodes, traffic=traffic, trace=trace)

    # run simulation for 500 ticks
    network.run(until=500)

    print(network.stats_summary)

if __name__ == '__main__':
    exit(main())