#!/usr/bin/env python3
#
# Simple benchmark script that tests NetworkTables throughput.
#
# To run the server
#
#     python3 benchmark.py
# 
# To run the client:
#
#    python3 benchmark.py 127.0.0.1
#
# In theory, the limiting factor is internal buffering and not network
# bandwidth, so running this test on a local machine should be mostly
# valid.
#
# On OSX with default write flush period, I average ~18hz
#

from __future__ import print_function

from argparse import ArgumentParser
import time

import logging
logging.basicConfig(level=logging.DEBUG)

from networktables import NetworkTables

class Benchmark(object):

    def __init__(self):
        
        parser = ArgumentParser()
        parser.add_argument('client', nargs='?', default=None)
        parser.add_argument('-r', '--rate', default=None, type=float)
        parser.add_argument('--send', default=False, action='store_true')
        
        args = parser.parse_args()
        
        if args.client:
            NetworkTables.initialize(server=args.client)
        else:
            NetworkTables.initialize()
         
        # Default write flush is 0.05, could adjust for less latency
        if args.rate is not None:
            print("Setting rate to %s" % args.rate)
            NetworkTables.setUpdateRate(args.rate)
        
        self.nt = NetworkTables.getTable('/benchmark')
        self.updates = 0
        
        self.nt.addTableListener(self.on_update)
        
        if args.send:
            self.send_benchmark()
        else:
            self.recv_benchmark()

    def on_update(self, *args):
        self.updates += 1

    def recv_benchmark(self):
        
        print("Starting to receive")
        
        last = None
        last_updates = 0
        
        while True:
            
            now = time.time()
            updates = self.updates
            
            if last is not None:
                rate = (updates - last_updates) / (now - last)
                print('Update rate:', rate)
            
            last = now
            last_updates = updates
            time.sleep(1)
            
    def send_benchmark(self):
        
        print("Sending")
        
        i = 0
        
        while True:
            i += 1
            self.nt.putNumber('key', i)
            time.sleep(0.0001)
    
if __name__ == '__main__':
    Benchmark()