# Oscilloscope
#
# Copyright 2008-2009 Red Hat, Inc.
#
# Arnaldo Carvalho de Melo <acme@redhat.com>
#
# Please check the tuna repository at:
# http://git.kernel.org/?p=linux/kernel/git/acme/tuna.git;a=tree
# For newer versions and to see it integrated with tuna
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation;
# version 2.1 of the License.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USA

import os
import sys
import gi
gi.require_version("Gtk", "3.0")
from gi.repository import Gtk
from gi.repository import Gdk
from gi.repository import GObject
from matplotlib.backends.backend_gtk3agg import \
    FigureCanvasGTK3Agg as figure_canvas
#from matplotlib.backends.backend_agg import \
#    FigureCanvasAgg as figure_canvas
import matplotlib.figure
import matplotlib.ticker
import numpy

class histogram_frame(Gtk.Frame):
    def __init__(self, title="Statistics", width=780, height=100,
                 max_value=500, nr_entries=10, facecolor="white"):

        Gtk.Frame.__init__(self)
        self.set_label(title)

        self.fraction = int(max_value / nr_entries)
        if self.fraction == 0:
            self.fraction = max_value
            nr_entries = 1
        self.max_value = max_value
        self.nr_entries = nr_entries
        self.nr_samples = 0

        table = Gtk.Table(3, self.nr_entries + 1, False)
        table.set_border_width(5)
        table.set_row_spacings(5)
        table.set_col_spacings(10)
        self.add(table)
        self.buckets = [0, ] * (nr_entries + 1)
        self.buckets_bar = [None, ] * (nr_entries + 1)
        self.buckets_counter = [None, ] * (nr_entries + 1)

        prefix = "<="
        for bucket in range(self.nr_entries + 1):
            bucket_range = (bucket + 1) * self.fraction
            if bucket_range > self.max_value:
                prefix = ">"
                bucket_range = self.max_value

            label = Gtk.Label("%s %d" % (prefix, bucket_range))
            label.set_alignment(0, 1)
            table.attach(label, 0, 1, bucket, bucket + 1, 0, 0, 0, 0)

            self.buckets_bar[bucket] = Gtk.ProgressBar()
            table.attach(self.buckets_bar[bucket], 1, 2, bucket, bucket + 1, 0, 0, 0, 0)

            self.buckets_counter[bucket] = Gtk.Label(label="0")
            label.set_alignment(0, 1)
            table.attach(self.buckets_counter[bucket], 2, 3, bucket, bucket + 1, 0, 0, 0, 0)

        self.modify_bg(Gtk.StateType.NORMAL, Gdk.color_parse(facecolor))

    def add_sample(self, sample):
        if sample > self.max_value:
            bucket = self.nr_entries
        else:
            bucket = int(sample / self.fraction)
        self.nr_samples += 1
        self.buckets[bucket] += 1

    def refresh(self):
        for bucket in range(self.nr_entries + 1):
            self.buckets_counter[bucket].set_text(str(self.buckets[bucket]))
            fraction = float(self.buckets[bucket]) / self.nr_samples
            self.buckets_bar[bucket].set_fraction(fraction)

    def reset(self):
        self.buckets = [0, ] * (self.nr_entries + 1)
        self.nr_samples = 0

class oscilloscope_frame(Gtk.Frame):

    def __init__(self, title="Osciloscope", width=780, height=360,
                 nr_samples_on_screen=250, graph_type='-',
                 max_value=500, plot_color="lightgreen",
                 bg_color="darkgreen", facecolor="white",
                 ylabel="Latency", picker=None):

        Gtk.Frame.__init__(self)
        self.set_label(title)

        self.font = {'fontname'   : 'Liberation Sans',
                     'color'      : 'b',
                     'fontweight' : 'bold',
                     'fontsize'   : 10}

        self.max_value = max_value
        self.nr_samples_on_screen = nr_samples_on_screen
        self.ind = numpy.arange(nr_samples_on_screen)
        self.samples = [0.0] * nr_samples_on_screen

        figure = matplotlib.figure.Figure(figsize=(10, 4), dpi=100,
                                          facecolor=facecolor)
        ax = figure.add_subplot(111)
        self.ax = ax
        ax.set_facecolor(bg_color)

        self.on_screen_samples = ax.plot(self.ind, self.samples, graph_type,
                                         color=plot_color, picker=picker)

        ax.set_ylim(0, max_value)
        ax.set_ylabel(ylabel, self.font)
        ax.set_xlabel("%d samples" % nr_samples_on_screen, self.font)
        ax.set_xticklabels([])
        ax.grid(True)

        for label in ax.get_yticklabels():
            label.set(fontsize=8)

        self.canvas = figure_canvas(figure)  # a Gtk.DrawingArea
        self.canvas.set_size_request(width, height)

        self.add(self.canvas)
        self.modify_bg(Gtk.StateType.NORMAL, Gdk.color_parse(facecolor))
        self.nr_samples = 0

    def add_sample(self, sample):
        del self.samples[0]
        self.samples.append(sample)
        self.on_screen_samples[0].set_data(self.ind, self.samples)
        self.nr_samples += 1
        if self.nr_samples <= self.nr_samples_on_screen:
            self.ax.set_xlabel("%d samples" % self.nr_samples, self.font)

    def reset(self):
        self.samples = [0.0] * self.nr_samples_on_screen
        self.nr_samples = 0
        self.on_screen_samples[0].set_data(self.ind, self.samples)

    def refresh(self):
        self.canvas.draw()
        return

def add_table_row(table, row, label_text, label_value="0"):
    label = Gtk.Label(label=label_text)
    label.set_use_underline(True)
    label.set_alignment(0, 1)
    table.attach(label, 0, 1, row, row + 1, 0, 0, 0, 0)

    label = Gtk.Label(label=label_value)
    table.attach(label, 1, 2, row, row + 1, 0, 0, 0, 0)
    return label

class system_info_frame(Gtk.Frame):
    def __init__(self, title="System", facecolor="white"):
        Gtk.Frame.__init__(self)
        self.set_label(title)

        self.modify_bg(Gtk.StateType.NORMAL, Gdk.color_parse(facecolor))

        table = Gtk.Table(3, 2, False)
        table.set_border_width(5)
        table.set_row_spacings(5)
        table.set_col_spacings(10)
        self.add(table)

        u = os.uname()
        add_table_row(table, 0, "Kernel Release", u[2])
        add_table_row(table, 1, "Architecture", u[4])
        add_table_row(table, 2, "Machine", u[1])

class oscilloscope(Gtk.Window):

    def __init__(self, get_sample=None, width=800, height=500,
                 nr_samples_on_screen=250,
                 graph_type='-', title="Osciloscope",
                 max_value=500, plot_color="lightgreen",
                 bg_color="darkgreen", facecolor="white",
                 ylabel="Latency",
                 picker=None,
                 snapshot_samples=0,
                 geometry=None, scale=True):

        Gtk.Window.__init__(self)
        if geometry:
            self.parse_geometry(geometry)
            width, height = self.get_size()
        else:
            self.set_default_size(width, height)

        self.get_sample = get_sample
        self.max_value = max_value
        self.snapshot_samples = snapshot_samples
        self.scale = scale

        self.set_title(title)

        vbox = Gtk.VBox()
        vbox.set_border_width(8)
        self.add(vbox)

        stats_frame = Gtk.Frame()
        stats_frame.set_label("Statistics")
        stats_frame.modify_bg(Gtk.StateType.NORMAL, Gdk.color_parse(facecolor))

        table = Gtk.Table(3, 2, False)
        table.set_border_width(5)
        table.set_row_spacings(5)
        table.set_col_spacings(10)
        stats_frame.add(table)

        self.min_label = add_table_row(table, 0, "Min")
        self.avg_label = add_table_row(table, 1, "Avg")
        self.max_label = add_table_row(table, 2, "Max")

        help_frame = Gtk.Frame()
        help_frame.set_label("Help")
        help_frame.modify_bg(Gtk.StateType.NORMAL, Gdk.color_parse(facecolor))

        table = Gtk.Table(4, 2, False)
        table.set_border_width(5)
        table.set_row_spacings(5)
        table.set_col_spacings(10)
        help_frame.add(table)

        add_table_row(table, 0, "Space", "Pause")
        add_table_row(table, 1, "S", "Snapshot")
        add_table_row(table, 2, "R", "Reset")
        add_table_row(table, 3, "Q", "Quit")

        self.scope = oscilloscope_frame("Scope",
                                        int(width * 0.94),
                                        int(height * 0.64),
                                        nr_samples_on_screen,
                                        max_value=max_value,
                                        graph_type=graph_type,
                                        picker=picker,
                                        ylabel=ylabel)

        self.hist = histogram_frame("Histogram", 0, 0, nr_entries=5,
                                    max_value=max_value)

        info_frame = system_info_frame()

        vbox_help_info = Gtk.VBox()
        vbox_help_info.pack_start(info_frame, False, False, 0)
        vbox_help_info.pack_end(help_frame, False, False, 0)
        hbox = Gtk.HBox()
        hbox.pack_start(vbox_help_info, False, False, 0)
        hbox.pack_start(stats_frame, False, False, 0)
        hbox.pack_end(self.hist, True, True, 0)

        vbox.pack_start(self.scope, True, True, 0)
        vbox.pack_end(hbox, True, False, 0)

        self.show_all()

        self.getting_samples = False
        self.refreshing_screen = False
        self.max = self.min = None
        self.avg = 0

    def add_sample(self, sample):
        if not self.max or self.max < sample:
            self.max = sample

        if not self.min or self.min > sample:
            self.min = sample

        self.avg = (self.avg + sample) / 2
        self.scope.add_sample(sample)
        self.hist.add_sample(sample)

    def refresh(self):
        if self.scale and self.max > self.scope.max_value:
            self.scope.max_value *= 2
            self.scope.ax.set_ylim(0, self.scope.max_value)
        self.scope.refresh()
        self.hist.refresh()
        while Gtk.events_pending():
            Gtk.main_iteration()

    def get_samples(self, fd, condition):
        try:
            sample = self.get_sample()
            prev_min, prev_avg, prev_max = self.min, self.avg, self.max

            self.add_sample(sample)

            if self.refreshing_screen:
                if self.min != prev_min:
                    self.min_label.set_text("%-6.3f" % self.min)
                if self.avg != prev_avg:
                    self.avg_label.set_text("%-6.3f" % self.avg)
                if self.max != prev_max:
                    self.max_label.set_text("%-6.3f" % self.max)

                self.refresh()

            if self.snapshot_samples == self.scope.nr_samples:
                self.snapshot()
                Gtk.main_quit()
        except:
            print("invalid sample, check the input format")
            pass
        return self.getting_samples

    def run(self, fd):
        self.connect("key_press_event", self.key_press_event)
        self.getting_samples = True
        self.refreshing_screen = True
        GObject.io_add_watch(fd, GObject.IO_IN | GObject.IO_PRI,
                             self.get_samples)

    def freeze_screen(self, state=False):
        self.refreshing_screen = state

    def stop(self):
        self.getting_samples = False
        self.refreshing_screen = False

    def snapshot(self):
        self.scope.canvas.print_figure("scope_snapshot.svg")

    def reset(self):
        self.scope.max_value = self.max_value
        self.scope.ax.set_ylim(0, self.scope.max_value)
        self.scope.reset()
        self.hist.reset()
        self.min = self.max_value
        self.max = 0
        self.avg = 0

    def key_press_event(self, widget, event):
        if event.keyval == ord(' '):
            self.freeze_screen(not self.refreshing_screen)
        elif event.keyval in (ord('s'), ord('S')):
            self.snapshot()
        elif event.keyval in (ord('r'), ord('R')):
            self.reset()
        elif event.keyval in (ord('q'), ord('Q')):
            Gtk.main_quit()

class ftrace_window(Gtk.Window):

    (COL_FUNCTION, ) = list(range(1))

    def __init__(self, trace, parent=None):
        Gtk.Window.__init__(self)
        try:
            self.set_screen(parent.get_screen())
        except AttributeError:
            self.connect('destroy', lambda *w: Gtk.main_quit())

        self.set_border_width(8)
        self.set_default_size(350, 500)
        self.set_title("ftrace")

        vbox = Gtk.VBox(False, 8)
        self.add(vbox)

        sw = Gtk.ScrolledWindow()
        sw.set_shadow_type(Gtk.ShadowType.ETCHED_IN)
        sw.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
        vbox.pack_start(sw, True, True)

        store = Gtk.ListStore(GObject.TYPE_STRING)

        for entry in trace:
            if entry[0] in ["#", "\n"] or entry[:4] == "vim:":
                continue
            iter = store.append()
            store.set(iter, self.COL_FUNCTION, entry.strip())

        treeview = Gtk.TreeView(store)
        treeview.set_rules_hint(True)

        column = Gtk.TreeViewColumn("Function", Gtk.CellRendererText(),
                                    text=self.COL_FUNCTION)
        treeview.append_column(column)

        sw.add(treeview)
        self.show_all()

class cyclictestoscope(oscilloscope):
    def __init__(self, max_value, snapshot_samples=0, nr_samples_on_screen=500,
                 delimiter=':', field=2, ylabel="Latency",
                 geometry=None, scale=True, sample_multiplier=1):
        oscilloscope.__init__(self, self.get_sample,
                              title="CyclictestoSCOPE",
                              nr_samples_on_screen=nr_samples_on_screen,
                              width=900, max_value=max_value,
                              picker=self.scope_picker,
                              snapshot_samples=snapshot_samples,
                              ylabel=ylabel, geometry=geometry,
                              scale=scale)

        self.connect("destroy", self.quit)
        self.delimiter = delimiter
        self.sample_multiplier = sample_multiplier
        self.field = field
        self.latency_tracer = os.access("/sys/kernel/debug/tracing/trace", os.R_OK)
        if self.latency_tracer:
            self.traces = [None,] * nr_samples_on_screen

    def scope_picker(self, line, mouseevent):
        if (not self.latency_tracer) or mouseevent.xdata is None:
            return False, dict()

        x = int(mouseevent.xdata)
        if self.traces[x]:
            fw = ftrace_window(self.traces[x], self)
        return False, dict()

    def get_sample(self):
        fields = sys.stdin.readline().split(self.delimiter)
        try:
            sample = float(fields[self.field]) * self.sample_multiplier
        except:
            print("fields=%s, self.field=%s,self.delimiter=%s" % (fields, self.field, self.delimiter))
            return None

        if self.latency_tracer:
            del self.traces[0]
            if sample > self.avg:
                print(sample)
                try:
                    f = open("/sys/kernel/debug/tracing/trace")
                    trace = f.readlines()
                    f.close()
                    f = open("/sys/kernel/debug/tracing/tracing_max_latency", "w")
                    f.write("0\n")
                    f.close()
                except:
                    trace = None
            else:
                print("-")
                trace = None

            self.traces.append(trace)

        return sample

    def run(self):
        oscilloscope.run(self, sys.stdin.fileno())

    def quit(self, x):
        Gtk.main_quit()
