#!/usr/bin/env python
# 
# ntpshmviz - graph the drift of NTP servers
# Written by Keane Wolter <daemoneye2@gmail.com>
#
# To do:
#
# 1. Add exit button so the user does not need to do <Ctrl>-w
#

import sys
# need numpy for float128, normal python floats are too small to hold a timespec
import numpy
try:
    import matplotlib.pyplot as PLT
    from matplotlib.figure import Figure
except:
    print("Please make sure matplotlib is installed properly:", sys.exc_info()[0])
    sys.exit(1)

class ntpOffset:
    def __init__(self, stream):
        # get the data
        self.read_data(stream)

        # display the data
        self.display()

    def display(self):
        # display the graphs

        # Alert the user that closing the graphs can be done with "Ctrl-w"
        print ("Please note that the graph can be closed with the key combination of <Ctrl-w>")

        PLT.figure()
        subplot_value = 211
        for ntp_item in self.ntp_data:
            # create the subplot for the data
            PLT.subplot(subplot_value)
            # setup and create the vlines graph
            t = numpy.arange(0, self.line_counts[ntp_item], 1)
            PLT.vlines(t, 0, self.ntp_data[ntp_item], color='r')
            # add labels
            PLT.title("NTP drift for " + ntp_item)
            PLT.xlabel('sample')
            PLT.ylabel('drift')
            # increment the subplot by 1.
            # this allows for each data group to have it's own graph
            subplot_value += 1
        # make sure there is no graph or data overlapping each other and display the graph
        PLT.tight_layout()
        PLT.show()

    def read_data(self, stream):
        # Reads data from a ntp log file.  Layout is:
        #
        #   - The keyword "sample"
        #   - The NTP unit from which it was collected.
        #   - Collection time of day, expressed in seconds
        #   - Receiver time of day, expressed in seconds
        #   - Clock time of day, expressed in seconds
        #   - Leep-second notification status
        #   - Source precision (log(2) of source jitter)

        self.ntp_data    = {} # data sets for each ntp unit
        self.line_counts = {} # Count of total lines for each ntp unit
        record           = [] # A single record in the file or data stream
        offset           = 0  # offset value to add to the array 
        self.lines       = 0  # width of graph

        for line in stream:
            if len(line.split(' ')) > 6:
                if line[:1] != '#':
                    line = line.lstrip()
                    record = line.split(' ')
                    try:
                        offset = numpy.float128(record[3]) - numpy.float128(record[4])
                    except:
                        print ("Invalid data: ", sys.exc_info()[0], ". Data was: ", line)
                        continue

                    # If the NTP unit is in the dictionary
                    # append the offset to the list
                    # Otherwise, create a new list in the dictionary
                    # and add the offset.
                    if record[1] not in self.ntp_data:
                        self.ntp_data[record[1]]  = []
                        self.line_counts[record[1]] = 0

                    self.ntp_data[record[1]].append(offset)
                    self.line_counts[record[1]] += 1
        stream.close()

if __name__ == "__main__":
    if (len(sys.argv) == 2):
        if (sys.argv[1] == "-V" or sys.argv[1] == "--version"):
            print ("Version 2.0")
            sys.exit()

        if (sys.argv[1] == "-h" or sys.argv[1] == "--help"):
            print ("Usage: <input stream> | ntpshmviz")
            print ("Example: cat SAMPLES | ntpshmviz")
            print ("")
            print ("Options:")
            print ("\t-h: help message")
            print ("\t-V: version information")
            print ("")
            print ("Report ntpshmviz bugs to Keane Wolter, daemoneye2@gmail.com")
            print ("gpsd homepage: www.catb.org/gpsd/")
            sys.exit()

    ntpOffset(sys.stdin)
    sys.exit()