#!/usr/bin/env python
#
# This file is Copyright (c) 2010 by the GPSD project
# SPDX-License-Identifier: BSD-2-clause
#
# Collect and plot latency-profiling data from a running gpsd.
# Requires gnuplot.
#
# This code runs compatibly under Python 2 and 3.x for x >= 2.
# Preserve this property!
from __future__ import absolute_import, print_function, division

import copy
import getopt
import math
import os
import signal
import socket
import sys
import time

import gps


class Baton(object):
    "Ship progress indication to stderr."

    def __init__(self, prompt, endmsg=None):
        self.stream = sys.stderr
        self.stream.write(prompt + "...")
        if os.isatty(self.stream.fileno()):
            self.stream.write(" \b")
        self.stream.flush()
        self.count = 0
        self.endmsg = endmsg
        self.time = time.time()
        return

    def twirl(self, ch=None):
        if self.stream is None:
            return
        if ch:
            self.stream.write(ch)
        elif os.isatty(self.stream.fileno()):
            self.stream.write("-/|\\"[self.count % 4])
            self.stream.write("\b")
        self.count = self.count + 1
        self.stream.flush()
        return

    def end(self, msg=None):
        if msg is None:
            msg = self.endmsg
        if self.stream:
            self.stream.write("...(%2.2f sec) %s.\n"
                              % (time.time() - self.time, msg))
        return


class plotter(object):
    "Generic class for gathering and plotting sensor statistics."

    def __init__(self):
        self.fixes = []
        self.start_time = int(time.time())
        self.watch = set(['TPV'])

    def whatami(self):
        "How do we identify this plotting run?"
        desc = "%s, %s, " % \
               (gps.misc.isotime(self.start_time),
                self.device.get('driver', "unknown"))
        if 'bps' in self.device:
            desc += "%d %dN%d, cycle %ds" % \
                (self.device['bps'], 9 - self.device['stopbits'],
                 self.device['stopbits'], self.device['cycle'])
        else:
            desc += self.device['path']
        return desc

    def collect(self, verbose, logfp=None):
        "Collect data from the GPS."
        try:
            self.session = gps.gps(host=host, port=port, verbose=verbose)
        except socket.error:
            sys.stderr.write("gpsprof: gpsd unreachable.\n")
            sys.exit(1)
        # Initialize
        self.session.read()
        if self.session.version is None:
            sys.stderr.write("gpsprof: requires gpsd to speak new protocol.\n")
            sys.exit(1)
        # Set parameters
        flags = gps.WATCH_ENABLE | gps.WATCH_JSON
        if self.requires_time:
            flags |= gps.WATCH_TIMING
        if device:
            flags |= gps.WATCH_DEVICE
        try:
            signal.signal(signal.SIGUSR1,
                          lambda empty, unused: sys.stderr.write(
                              "%d of %d (%d%%)..."
                              % (await - countdown, await,
                                 ((await - countdown) * 100.0 / await))))
            signal.siginterrupt(signal.SIGUSR1, False)
            self.session.stream(flags, device)
            baton = Baton("gpsprof: %d looking for fix" % os.getpid(), "done")
            countdown = await
            basetime = time.time()
            while countdown > 0:
                if self.session.read() == -1:
                    sys.stderr.write("gpsprof: gpsd has vanished.\n")
                    sys.exit(1)
                baton.twirl()
                if self.session.data["class"] == "ERROR":
                    sys.stderr.write(" ERROR: %s.\n"
                                     % self.session.data["message"])
                    sys.exit(1)
                if self.session.data["class"] == "DEVICES":
                    if len(self.session.data["devices"]) != 1 and not device:
                        sys.stderr.write("ERROR: multiple devices connected, "
                                         "you must explicitly specify the "
                                         "device.\n")
                        sys.exit(1)
                    for i in range(len(self.session.data["devices"])):
                        self.device = copy.copy(
                            self.session.data["devices"][i])
                        if self.device['path'] == device:
                            break
                if self.session.data["class"] == "WATCH":
                    if ((self.requires_time
                         and not self.session.data.get("timing"))):
                        sys.stderr.write("timing is not enabled.\n")
                        sys.exit(1)
                # Log before filtering - might be good for post-analysis.
                if logfp:
                    logfp.write(self.session.response)
                # Ignore everything but what we're told to
                if self.session.data["class"] not in self.watch:
                    continue
                # We can get some funky artifacts at start of self.session
                # apparently due to RS232 buffering effects. Ignore
                # them.
                if ((threshold and
                     time.time() - basetime < self.session.cycle * threshold)):
                    continue
                if self.session.fix.mode <= gps.MODE_NO_FIX:
                    continue
                if self.sample():
                    if countdown == await:
                        sys.stderr.write("first fix in %.2fsec, gathering %d "
                                         "samples..."
                                         % (time.time() - basetime, await))
                    countdown -= 1
            baton.end()
        finally:
            self.session.stream(gps.WATCH_DISABLE | gps.WATCH_TIMING)
            signal.signal(signal.SIGUSR1, signal.SIG_DFL)

    def replot(self, infp):
        "Replot from a JSON log file."
        baton = Baton("gpsprof: replotting", "done")
        self.session = gps.gps(host=None)
        for line in infp:
            baton.twirl()
            self.session.unpack(line)
            if self.session.data["class"] == "DEVICES":
                self.device = copy.copy(self.session.data["devices"][0])
            elif self.session.data["class"] not in self.watch:
                continue
            self.sample()
        baton.end()

    def dump(self):
        "Dump the raw data for post-analysis."
        return self.header() + self.data()


class spaceplot(plotter):
    "Spatial scattergram of fixes."
    name = "space"
    requires_time = False

    def __init__(self):
        plotter.__init__(self)
        self.recentered = []

    def d(self, a, b):
        return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)

    def sample(self):
        # Watch out for the NaN value from gps.py.
        if (self.session.valid &
            (gps.ONLINE_SET | gps.PACKET_SET | gps.LATLON_SET)):
            self.fixes.append((self.session.fix.latitude,
                               self.session.fix.longitude,
                               self.session.fix.altitude))
        return True

    def header(self):
        return "# Position uncertainty, %s\n" % self.whatami()

    def postprocess(self):
        if not self.recentered:
            # centroid is just arithmetic avg of lat,lon
            self.centroid = (sum([x[0] for x in self.fixes]) / len(self.fixes),
                             sum([x[1] for x in self.fixes]) / len(self.fixes))
            # Sort fixes by distance from centroid
            self.fixes.sort(key=lambda p: self.d(self.centroid, p))
            # Convert fixes to offsets from centroid in meters
            self.recentered = [
                gps.MeterOffset(self.centroid, fix[:2]) for fix in self.fixes]

    def data(self):
        res = ""
        for i in range(len(self.recentered)):
            (lat, lon) = self.recentered[i][:2]
            (raw1, raw2, alt) = self.fixes[i]
            res += "%.9f\t%.9f\t%.9f\t%.9f\t%.9f\n" \
                % (lat, lon, raw1, raw2, alt)
        return res

    def plot(self):
        # Compute CEP(50%)
        cep_meters = gps.misc.EarthDistance(
            self.centroid[:2], self.fixes[int(len(self.fixes) * 0.50)][:2])
        cep95_meters = gps.misc.EarthDistance(
            self.centroid[:2], self.fixes[int(len(self.fixes) * 0.95)][:2])
        cep99_meters = gps.misc.EarthDistance(
            self.centroid[:2], self.fixes[int(len(self.fixes) * 0.99)][:2])
        alt_sum = 0.0
        alt_num = 0
        alt_fixes = []
        lon_max = -9999
        for i in range(len(self.recentered)):
            (_lat, lon) = self.recentered[i][:2]
            (_raw1, _raw2, alt) = self.fixes[i]
            if not gps.isnan(alt):
                alt_sum += alt
                alt_fixes.append(alt)
                alt_num += 1
            if lon > lon_max:
                lon_max = lon
        if alt_num == 0:
            alt_avg = gps.NaN
            alt_ep = gps.NaN
        else:
            alt_avg = alt_sum / alt_num
            # Sort fixes by distance from average altitude
            alt_fixes.sort(key=lambda a: abs(alt_avg - a))
            alt_ep = abs(alt_fixes[len(alt_fixes) // 2] - alt_avg)
        if self.centroid[0] < 0:
            latstring = "%.9fS" % -self.centroid[0]
        elif self.centroid[0] == 0:
            latstring = "0"
        else:
            latstring = "%.9fN" % self.centroid[0]
        if self.centroid[1] < 0:
            lonstring = "%.9fW" % -self.centroid[1]
        elif self.centroid[1] == 0:
            lonstring = "0"
        else:
            lonstring = "%.9fE" % self.centroid[1]
        fmt = "set autoscale\n"
        fmt += "set format x \"%.3f\"\n"
        fmt += "set format y \"%.3f\"\n"
        fmt += 'set key below\n'
        fmt += 'set key title "%s"\n' % gps.misc.isotime(int(time.time()))
        fmt += 'set size ratio -1\n'
        fmt += 'set style line 2 pt 1\n'
        fmt += 'set style line 3 pt 2\n'
        fmt += 'set xlabel "Meters east from %s"\n' % lonstring
        fmt += 'set xtic rotate by -45\n'
        fmt += 'set ylabel "Meters north from %s"\n' % latstring
        fmt += 'set border 15\n'
        if not gps.isnan(alt_avg):
            fmt += 'set y2label "Meters Altitude from %.3f"\n' % alt_avg
            fmt += 'set ytics nomirror\n'
            fmt += 'set y2tics\n'
            fmt += "set format y2 \"%.3f\"\n"
        fmt += 'cep=%.9f\n' \
            % self.d((0, 0), self.recentered[len(self.fixes) // 2])
        fmt += 'cep95=%.9f\n' \
            % self.d((0, 0), self.recentered[int(len(self.fixes) * 0.95)])
        fmt += 'cep99=%.9f\n' \
            % self.d((0, 0), self.recentered[int(len(self.fixes) * 0.99)])
        fmt += 'set parametric\n'
        fmt += 'set trange [0:2*pi]\n'
        fmt += 'cx(t, r) = sin(t)*r\n'
        fmt += 'cy(t, r) = cos(t)*r\n'
        fmt += 'chlen = cep/20\n'
        fmt += "set arrow from -chlen,0 to chlen,0 nohead\n"
        fmt += "set arrow from 0,-chlen to 0,chlen nohead\n"
        if len(self.fixes) > 1000:
            plot_style = 'dots'
        else:
            plot_style = 'points'
        fmt += 'plot "-" using 1:2 with ' + plot_style \
               + ' ls 3 title "%d GPS fixes" ' % (len(self.fixes))
        if not gps.isnan(alt_avg):
            fmt += ', "-" using ( %.3f ):($5 < 100000 ? $5 - %.3f : 1/0)' \
                   'axes x1y2 with %s ls 2 title ' \
                   '" %d Altitude fixes, Average = %.3f, EP (50%%) = %.3f"' \
                   % (lon_max + 1, alt_avg, plot_style, alt_num, alt_avg,
                      alt_ep)
        fmt += ', cx(t, cep),cy(t, cep) ls 1 title "CEP (50%%) = %.3f meters"' \
            % (cep_meters)
        fmt += ', cx(t, cep95),cy(t, cep95) title "CEP (95%%) = %.3f meters"' \
            % (cep95_meters)
        fmt += ', cx(t, cep99),cy(t, cep99) title "CEP (99%%) = %.3f meters"' \
            % (cep99_meters)
        fmt += "\n"
        fmt += self.header()
        fmt += self.data()
        if not gps.isnan(alt_avg):
            fmt += "e\n" + self.data()
        return fmt


class timeplot(plotter):
    "Time drift against PPS."
    name = "time"
    requires_time = True

    def __init__(self):
        plotter.__init__(self)
        self.watch = set(['PPS'])

    def sample(self):
        if self.session.data["class"] == "PPS":
            self.fixes.append((self.session.data['real_sec'],
                               self.session.data['real_nsec'],
                               self.session.data['clock_sec'],
                               self.session.data['clock_nsec']))
        return True

    def header(self):
        return "# Time drift against PPS, %s\n" % self.whatami()

    def postprocess(self):
        pass

    def data(self):
        res = ""
        for (real_sec, real_nsec, clock_sec, clock_nsec) in self.fixes:
            res += "%d\t%d\t%d\t%d\n" % (real_sec, real_nsec, clock_sec,
                                         clock_nsec)
        return res

    def plot(self):
        fmt = '''\
set autoscale
set key below
set ylabel "System clock delta from GPS time (nsec)"
plot "-" using 0:((column(1)-column(3))*1e9 + (column(2)-column(4))) \
 title "Delta" with impulses
'''
        return fmt + self.header() + self.data()


class uninstrumented(plotter):
    "Total times without instrumentation."
    name = "uninstrumented"
    requires_time = False

    def __init__(self):
        plotter.__init__(self)

    def sample(self):
        if self.session.fix.time:
            seconds = time.time() - gps.misc.isotime(self.session.data.time)
            self.fixes.append(seconds)
            return True
        else:
            return False

    def header(self):
        return "# Uninstrumented total latency, " + self.whatami() + "\n"

    def postprocess(self):
        pass

    def data(self):
        res = ""
        for seconds in self.fixes:
            res += "%2.6lf\n" % seconds
        return res

    def plot(self):
        fmt = '''\
set autoscale
set key below
set key title "Uninstrumented total latency"
plot "-" using 0:1 title "Total time" with impulses
'''
        return fmt + self.header() + self.data()


class instrumented(plotter):
    "Latency as analyzed by instrumentation."
    name = "instrumented"
    requires_time = True

    def __init__(self):
        plotter.__init__(self)

    def sample(self):
        if 'rtime' in self.session.data:
            self.fixes.append((gps.misc.isotime(self.session.data['time']),
                               self.session.data["chars"],
                               self.session.data['sats'],
                               self.session.data['sor'],
                               self.session.data['rtime'],
                               time.time()))
            return True
        else:
            return False

    def header(self):
        res = "# Analyzed latency, " + self.whatami() + "\n"
        res += "#-- Fix time --  - Chars -  --   Latency  - RS232-  " \
               "Analysis  - Recv -\n"
        return res

    def postprocess(self):
        pass

    def data(self):
        res = ""
        for (fix_time, chars, sats, start, xmit, recv) in self.fixes:
            rs232_time = (chars * 10.0) / self.device['bps']
            res += "%.3f  %9u  %2u  %.6f  %.6f  %.6f  %.6f\n" \
                % (fix_time, chars, sats, start - fix_time,
                   (start - fix_time) + rs232_time, xmit - fix_time,
                   recv - fix_time)
        return res

    def plot(self):
        legends = (
            "Reception delta",
            "Analysis time",
            "RS232 time",
            "Fix latency",
        )
        fmt = '''\
set autoscale
set key title "Analyzed latency"
set key below
plot \\\n'''
        for (i, legend) in enumerate(legends):
            j = len(legends) - i + 4
            fmt += '    "-" using 0:%d title "%s" with impulses, \\\n' \
                % (j, legend)
        fmt = fmt[:-4] + "\n"
        return fmt + self.header() + (self.data() + "e\n") * len(legends)

formatters = (spaceplot, timeplot, uninstrumented, instrumented)

if __name__ == '__main__':
    try:
        (options, arguments) = getopt.getopt(sys.argv[1:],
                                             "d:f:hl:m:n:rs:t:T:D:")

        plotmode = "space"
        raw = False
        title = None
        threshold = 0
        await = 100
        verbose = 0
        terminal = None
        dumpfile = None
        logfp = None
        redo = False
        for (switch, val) in options:
            if switch == '-f':
                plotmode = val
            elif switch == '-m':
                threshold = int(val)
            elif switch == '-n':
                if val[-1] == 'h':
                    await = int(val[:-1]) * 360
                else:
                    await = int(val)
            elif switch == '-t':
                title = val
            elif switch == '-T':
                terminal = val
            elif switch == '-d':
                dumpfile = val
            elif switch == '-l':
                logfp = open(val, "w")
            elif switch == '-r':
                redo = True
            elif switch == '-D':
                verbose = int(val)
            elif switch == '-h':
                sys.stderr.write(
                    "usage: gpsprof [-h] [-D debuglevel] [-m threshold] "
                    "[-n samplecount] [-d]\n"
                    + "\t[-f {" + "|".join([x.name for x in formatters])
                    + "}] [-s speed] [-t title] [-T terminal] "
                    "[server[:port[:device]]]\n")
                sys.exit(0)

        (host, port, device) = ("localhost", "2947", None)
        if len(arguments):
            args = arguments[0].split(":")
            if len(args) >= 1:
                host = args[0]
            if len(args) >= 2:
                port = args[1]
            if len(args) >= 3:
                device = args[2]

        # Select the plotting mode
        if plotmode:
            for formatter in formatters:
                if formatter.name == plotmode:
                    plot = formatter()
                    break
            else:
                sys.stderr.write("gpsprof: no such formatter.\n")
                sys.exit(1)
        # Get fix data from the GPS
        if redo:
            plot.replot(sys.stdin)
        else:
            plot.collect(verbose, logfp)
        plot.postprocess()
        # Save the timing data (only) for post-analysis if required.
        if dumpfile:
            with open(dumpfile, "w") as fp:
                fp.write(plot.dump())
        if logfp:
            logfp.close()
        # Ship the plot to standard output
        if not title:
            title = plot.whatami()
        if terminal:
            sys.stdout.write("set terminal %s size 800,600\n" % terminal)
        sys.stdout.write("set title \"%s\"\n" % title)
        sys.stdout.write(plot.plot())
    except KeyboardInterrupt:
        pass

# The following sets edit modes for GNU EMACS
# Local Variables:
# mode:python
# End: