#!/usr/bin/env python
# Plot the changes in /proc/interrupts over time
# Date: 2014-06-17
# Author: Peter Wu <peter@lekensteyn.nl>

# Wishlist:
# Thicker legend lines
# Nicer smoothing

import matplotlib.pyplot as plt
import matplotlib
from collections import deque, OrderedDict
import numpy as np
from scipy.interpolate import interp1d
import threading
from argparse import ArgumentParser

# Number of seconds to show in the graph
XRANGE = 60
# Delay between updating the graph
INTERVAL = .5
# Log scale base or 0 to disable logarithmic y scaling
LOG_SCALE_BASE = 10
# Whether to enable smooth curves or not
SMOOTH_CURVES = True

MARKER_DEFAULT = 'o'
MARKER_SELECTED = 'v'

# 26 colors from http://graphicdesign.stackexchange.com/a/3815
# "A Colour Alphabet and the Limits of Colour Coding"
COLORS = ['#F0A3FF', '#0075DC', '#993F00', '#4C005C', '#191919', '#005C31',
'#2BCE48', '#FFCC99', '#808080', '#94FFB5', '#8F7C00', '#9DCC00', '#C20088',
'#003380', '#FFA405', '#FFA8BB', '#426600', '#FF0010', '#5EF1F2', '#00998F',
'#E0FF66', '#740AFF', '#990000', '#FFFF80', '#FFFF00', '#FF5005']
# From alex440's comment (currently not used)
ALT_COLORS = ['#023FA5', '#7D87B9', '#BEC1D4', '#D6BCC0', '#BB7784', '#FFFFFF',
'#4A6FE3', '#8595E1', '#B5BBE3', '#E6AFB9', '#E07B91', '#D33F6A', '#11C638',
'#8DD593', '#C6DEC7', '#EAD3C6', '#F0B98D', '#EF9708', '#0FCFC0', '#9CDED6',
'#D5EAE7', '#F3E1EB', '#F6C4E1', '#F79CD4']

input_filename, output_filename = '', ''

parser = ArgumentParser(description='Monitor /proc/interrupts changes')
parser.add_argument('-i', dest='input_filename',
                    help='Source to read interrupts entries from')
parser.add_argument('-o', dest='output_filename',
                    help='Target file to write interrupt entries to')
parser.add_argument('-l', dest='limit_seconds', type=float,
                    help='Maximum number of seconds to remember in the graph (0 for unbounded)')
parser.add_argument('-x', '--xrange', type=float,
                    help='Range of x values to display')
parser.add_argument('-I', '--interval', type=float,
                    help='Seconds to wait between updates')
parser.add_argument('--log', type=float,
                    help='Log base of y axis or 0 to disable log y axis')
parser.add_argument('--no-smooth', action='store_false',
                    help='Whether to enable curve fitting')
args = parser.parse_args()
if args.input_filename:
    input_filename = args.input_filename
if args.output_filename:
    output_filename = args.output_filename
if args.xrange:
    XRANGE = args.xrange
if args.interval:
    INTERVAL = args.interval
if args.log is not None:
    LOG_SCALE_BASE = args.log
if args.no_smooth is not None:
    SMOOTH_CURVES = args.no_smooth
if args.limit_seconds is not None:
    if args.limit_seconds > 0:
        x_entries_max = int(args.limit_seconds / INTERVAL)
    else:
        # Unbounded
        x_entries_max = None
else:
    x_entries_max = int(XRANGE / INTERVAL)

def is_line_ok(name, yvalues):
    """Returns True if a line should be displayed for this name."""
    if max(yvalues) < 5:
        return False

    names_ok = ['hci', 'timer']
    for name_ok_part in names_ok:
        if name_ok_part in name:
            return True

    # Accept all
    return True

# Fix Unicode font
matplotlib.rc('font', family='DejaVu Sans')

# From http://stackoverflow.com/a/490090
def synchronized(lock):
    """Synchronization decorator."""
    def wrap(f):
        def newFunction(*args, **kw):
            with lock:
                return f(*args, **kw)
        return newFunction
    return wrap

if not input_filename:
    input_filename = '/proc/interrupts'
input_is_proc = input_filename == '/proc/interrupts'
input_file = open(input_filename)
output_file = open(output_filename, 'a') if output_filename else None
output_writer = output_file.write if output_file else None

def get_numbers():
    # TODO: may break the graph if a line disappears
    if input_is_proc:
        # Rewind stream
        input_file.seek(0)
        return parse_raw(input_file, output_writer)
    else:
        return parse_raw(input_file, output_writer)

def parse_raw(pi, line_callback=None):
    ncpus = None
    for line in pi:
        # Treat empty lines as boundary
        if not line.strip():
            break
        if line_callback:
            line_callback(line)

        if ncpus is None:
            ncpus = len(line.split())
            continue
        name, values = line.split(':', 1)
        name = name.strip()
        values = values.strip().split(None, ncpus)
        if len(values) >= ncpus:
            # Name is ID + description for uniqueness
            name += ':' + values[-1];
            yield name, sum(int(values[i]) for i in range(0, ncpus))

    # Signal end of entry
    if line_callback:
        line_callback('\n')

prev = OrderedDict()
def get_diffs():
    for name, n in get_numbers():
        if name in prev:
            yield name, n - prev[name]
        else:
            yield name, 0
        prev[name] = n

plt.ylabel(u'\u0394interrupts')
plt.xlabel(u'time (sec)')
plt.grid('on')
#plt.ion() # Not necessary if show() does not block.
plt.show(block=False)
if LOG_SCALE_BASE > 0:
    plt.yscale('log', nonpositive='clip', base=LOG_SCALE_BASE)

# Used when picking a new line or in update()
update_legend = False
# Lines depicting interrupts changes (the "real" markers and smoothed lines)
lines, smooth_lines = {}, {}
### BEGIN EVENTS

# After pressing ^W, stop the main loop
running = True
def on_close(event):
    global running
    running = False
plt.connect('close_event', on_close)

# Space toggles updating
paused = False
def on_keypress(event):
    global paused
    if event.key == ' ':
        paused = not paused
    update_title()
plt.connect('key_press_event', on_keypress)

last_selected, last_smooth = None, None
lines_lock = threading.Lock()
@synchronized(lines_lock)
def select_line(line):
    global last_selected, last_smooth
    # HACK: needed to support selection of legend items
    name = line.get_label()
    line = lines[name]
    if last_selected:
        last_selected.set_marker(MARKER_DEFAULT)
        for someline in (last_selected, last_smooth):
            if someline:
                someline.set_linewidth(someline.get_linewidth() / 4)
                someline.set_zorder(someline.get_zorder() - 1)
    if last_selected == line:
        # Same line selected again. Clicking acts as a toggle, so do not mark
        # the same line again.
        last_selected = None
    else:
        line.set_marker(MARKER_SELECTED)
        last_smooth = smooth_lines[name] if name in smooth_lines else None
        for someline in (line, last_smooth):
            if someline:
                someline.set_linewidth(someline.get_linewidth() * 4)
                someline.set_zorder(someline.get_zorder() + 1)
        last_selected = line

def on_pick(event):
    global update_legend
    artist = event.artist
    if isinstance(artist, matplotlib.lines.Line2D):
        select_line(artist)
        update_legend = True
        do_draw()
plt.connect('pick_event', on_pick)

### END EVENTS

names = [name for name, _ in get_diffs()]
yvalues = {}
# Initialize y values for each name
for name in names:
    ydata = deque([], x_entries_max)
    yvalues[name] = ydata

def update_title():
    title = input_filename
    if paused:
        title += ' (paused - press Space to resume)'
    plt.gcf().canvas.set_window_title(title)

# Lock to avoid updating the UI while the data is being refreshed
data_lock = threading.Lock()
start_time = 0
@synchronized(data_lock)
def refresh_data():
    global start_time
    # Update data
    updated = False
    for name, n in get_diffs():
        yvalues[name].append(n)
        updated = True
    if updated:
        start_time += INTERVAL
    return updated

last_min_x = 0
@synchronized(lines_lock)
@synchronized(data_lock)
def update():
    """Reads new data and updates the line values."""
    global update_legend, last_min_x
    # Update lines
    for name in names:
        ys = yvalues[name]
        # Consider only strictly positive values
        ydata = [y for y in ys if y > 0]
        # Skip time only if the size of yvalues is bounded
        skip_time = max(0, start_time - len(ys) * INTERVAL)
        xdata = [skip_time + i * INTERVAL for i, y in enumerate(ys) if y > 0]

        if ydata and is_line_ok(name, ys):
            # Data is significant, show it
            if name in lines:
                lines[name].set_data(xdata, ydata)
            else:
                color = COLORS[names.index(name) % len(COLORS)]
                lines[name], = plt.plot(xdata, ydata, MARKER_DEFAULT + '-',
                                        label=name,
                                        color=color)
                lines[name].set_pickradius(5) # Make selectable
                update_legend = True

            # Smooth curve
            ydata_len = len(ydata)
            if ydata_len > 3 and SMOOTH_CURVES:
                min_x = min(xdata)
                max_x = max(xdata)
                xnew = np.linspace(min_x, max_x, round(1 + max_x - min_x) * 8)
                #ynew = spline(xdata, ydata, xnew)
                # quadratic and cubic splines give too much deviations
                ynew = interp1d(xdata, ydata, kind='slinear')(xnew)
                if not name in smooth_lines:
                    line = lines[name]
                    smooth_lines[name], = plt.plot(xnew, ynew,
                            linewidth=line.get_linewidth(),
                            zorder=line.get_zorder(),
                            color=line.get_color())
                    # Smooth line is shown, hide straight lines
                    line.set_linestyle('')
                else:
                    smooth_lines[name].set_data(xnew, ynew)
            elif name in smooth_lines:
                smooth_lines[name].remove()
                del smooth_lines[name]
                # No smooth line is shown, fallback to straight lines
                lines[name].set_linestyle('-')
        elif name in lines:
            # Data is insignificant, remove previous line
            lines[name].remove()
            del lines[name]
            update_legend = True

    largest = 10
    for name in yvalues:
        ydata = yvalues[name]
        if is_line_ok(name, ydata):
            largest = max(largest, max(ydata))
    # Update iff graph becomes too large
    #ymin, ymax = plt.ylim()
    #if ymax - ymin < largest:
    #    plt.ylim(ymin, ymin + largest)
    plt.ylim(1, largest)

    # The first XRANGE items fit in the screen
    min_x = max(start_time - XRANGE, 0)
    cur_min_x, _ = plt.xlim()
    # If the current range is equal or further than the previously recorded
    # range, consider it sticky and update the xlim with the new range.
    if cur_min_x >= last_min_x:
        last_min_x = min_x
        if cur_min_x <= min_x:
            plt.xlim(min_x, min_x + XRANGE)

def do_draw():
    """Actually draw the graph, updating the legend if necessary."""
    global update_legend
    # update legend if a line gets added, changed or removed
    if update_legend:
        # TODO should probably create legend once and then use set_label to
        # update? Right now it gives a warning:
        # ./interrupts-graph.py:335: MatplotlibDeprecationWarning: Adding an
        # axes using the same arguments as a previous axes currently reuses the
        # earlier instance.  In a future version, a new instance will always be
        # created and returned.  Meanwhile, this warning can be suppressed, and
        # the future behavior ensured, by passing a unique label to each axes
        # instance.
        old_legend = plt.axes().get_legend()
        # Private API, use it to restore the legend position.
        old_loc = old_legend and old_legend._loc or 'upper left'

        legend = plt.legend(loc=old_loc,
                            framealpha=.5,
                            fontsize='small')
        legend.set_draggable(True)

        # Enable selecting a line by clicking in the legend
        for line in legend.get_lines():
            line.set_pickradius(5)

        update_legend = False

    plt.draw()

update_title()

# Separate worker for fetching data
t = None
def refresh_data_timer():
    global t
    t = threading.Timer(INTERVAL, refresh_data_timer)
    t.start()
    refresh_data()

if not input_is_proc:
    # Read all yvalues data from file
    while refresh_data():
        pass
    update()
    do_draw()
    # Block until exited
    plt.show()
else:
    refresh_data_timer()

    while running:
        if not paused:
            update()
        do_draw()
        plt.pause(INTERVAL)

    # Cancel any scheduled timer
    t.cancel()