#!/usr/bin/python

from PyQt4 import QtCore, QtGui

import subprocess, os, select, sys
import numpy as np
import scipy.signal as signal
import itertools, struct, warnings
import pickle
import h5py
import fcntl
import io

from vispy import gloo, app, scene, visuals
from vispy.visuals import transforms
from vispy.visuals.collections import SegmentCollection

sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from filtering import FilterPlotWindow

nrows = 8
ncols = 4
nchans = nrows*ncols

MICROVOLTS_PER_COUNT = 0.195

MIN_CHIP_NUMBER = 0
MAX_CHIP_NUMBER = 31    # Willows can have at most chips 0 - 31 active at once.
LABEL_HEIGHT = 60       # it looks fine for a line of text to be this tall

CHANGE_CHIP_MS = 500    # Aggregate chip change events before informing WillowGUI

# wait 10 seconds before giving up on receiving streaming data from proto2bytes
# and asking GUI to pause streaming in hardware
STREAM_GIVEUP_TIME_MS = 10000
# wait 100 ms before giving up on a single proto2bytes poll
POLL_TIME_MS = 100

class StreamWindow(QtGui.QWidget):

    def __init__(self, params, proto2bytes_filename):
        super(StreamWindow, self).__init__(None)

        if params['impedancePath'] != None:
            f = h5py.File(params['impedancePath'])
            self.impedance = f['impedanceMeasurements']
            self.display_impedance = True
        else:
            self.impedance = None
            self.display_impedance = False

        if params['probeMapPath'] != None:
            self.probeMap = pickle.load(open(params['probeMapPath'],
                                             'rb')) # (shank, col, row): channel
        else:
            self.probeMap = None

        ymin = -6390
        ymax = 6390
        self.deltay_uv = ymax - ymin
        self.refreshRate = params['refreshRate']
        self.deltax_ms = params['xrange']

        self.labels = QtGui.QWidget()
        self.infoLabel = QtGui.QLabel()
        self.labelLayout = QtGui.QHBoxLayout()
        self.labelLayout.addWidget(self.infoLabel)
        self.labels.setMaximumHeight(2*LABEL_HEIGHT)    # label is 2 lines tall
        if self.probeMap is not None:
            self.plotLabel = QtGui.QLabel()
            self.labels.setMaximumHeight(4*LABEL_HEIGHT)    # 2 more lines tall
            self.labelLayout.addWidget(self.plotLabel)
        self.labels.setLayout(self.labelLayout)

        #########################
        # vispy plotting stuff
        #########################

        sr = 30000  # sample rate
        frame_rate = self.refreshRate
        self.frame_period = 1000//frame_rate     # frame period, in ms
        self.nsamps = 30*self.deltax_ms   # number of samples to display
        self.nrefresh = sr//frame_rate   # new samples collected before refresh

        # apparently, the smoothest way to switch display between fullscreen and
        # not-fullscreen is just to do all the work both ways all the time.
        self.canvas_embedded = ChipCanvas(self.nsamps, self.nrefresh, ymin, ymax, self, parent=self)
        self.canvas_popout = ChipCanvas(self.nsamps, self.nrefresh, ymin, ymax, self, parent=None)

        self.initChip(params['chip'])
        self.canvas_embedded.native.show()
        # makes SEEKRIT SPEEDZ when opening fullscreen first time
        # this is mostly only helpful on crappy computers
        self.canvas_popout.native.showFullScreen()
        self.poppedIn = True

        self.plotBuff = np.zeros((nchans, self.nsamps), dtype='float')
        self.plotBuff_raw = np.zeros((nchans, self.nsamps), dtype='float')
        self.newBuff = np.zeros((nchans, self.nrefresh), dtype='float')
        self.draw_point = 0
        # nchans channels of 2-byte data, gathered over self.nrefresh samples
        size_data = 2*nchans*self.nrefresh
        self.new_data_ba = bytearray([0] * size_data)
        self.new_data_mv = memoryview(self.new_data_ba)


        self.hwif_request('setSubsamples_byChip', str(self.chip))

        self.p2b_filename = os.path.abspath(proto2bytes_filename)

        ###################
        # Top-level stuff
        ##################

        self.startButton = QtGui.QPushButton()
        self.startButton.setIcon(QtGui.QIcon('../img/play.png'))
        self.startButton.setIconSize(QtCore.QSize(25,25))
        self.startButton.clicked.connect(self.startStreaming)

        self.stopButton = QtGui.QPushButton()
        self.stopButton.setIcon(QtGui.QIcon('../img/pause.png'))
        self.stopButton.setIconSize(QtCore.QSize(25,25))
        self.stopButton.clicked.connect(self.stopStreaming)

        self.streaming = False

        self.filterCheckbox = QtGui.QCheckBox('Filtering displayed data')
        self.filterCheckbox.setCheckState(QtCore.Qt.Unchecked)
        self.filtering = False
        self.filterCheckbox.toggled.connect(self.toggleFiltering)
        self.plotHButton = QtGui.QPushButton('Plot Transfer Function')
        self.plotHButton.clicked.connect(self.launchFilterPlot)
        self.lowcutLabel = QtGui.QLabel('Lower cutoff (Hz)')
        self.lowcutLine = QtGui.QLineEdit('300')
        self.lowcutLine.setMaxLength(6)
        self.lowcutLine.setValidator(QtGui.QIntValidator())
        self.lowcutAccept = QtGui.QPushButton('Update lowcut')
        self.lowcutAccept.clicked.connect(self.setLowcut)
        self.setLowcut()
        self.highcutLabel = QtGui.QLabel('Higher cutoff (Hz)')
        self.highcutLine = QtGui.QLineEdit('9500')
        self.highcutLine.setMaxLength(6)
        self.highcutLine.setValidator(QtGui.QIntValidator())
        self.highcutLine.editingFinished.connect(self.setHighcut)
        self.highcutAccept = QtGui.QPushButton('Update highcut')
        self.highcutAccept.clicked.connect(self.setHighcut)
        self.setHighcut()


        self.filterPanel = QtGui.QWidget()
        filterAccept = QtGui.QWidget()
        filterParams = QtGui.QWidget()
        filterAcceptLayout = QtGui.QVBoxLayout()
        filterParamsLayout = QtGui.QGridLayout()
        filterAcceptLayout.addWidget(self.filterCheckbox)
        filterAcceptLayout.addWidget(self.plotHButton)
        filterParamsLayout.addWidget(self.lowcutLabel, 0,0)
        filterParamsLayout.addWidget(self.lowcutLine, 1,0)
        filterParamsLayout.addWidget(self.lowcutAccept, 2,0)
        filterParamsLayout.addWidget(self.highcutLabel, 0,1)
        filterParamsLayout.addWidget(self.highcutLine, 1,1)
        filterParamsLayout.addWidget(self.highcutAccept, 2,1)
        filterAccept.setLayout(filterAcceptLayout)
        filterParams.setLayout(filterParamsLayout)
        filterLayout = QtGui.QHBoxLayout()
        filterLayout.addWidget(filterAccept)
        filterLayout.addWidget(filterParams)
        self.filterPanel.setLayout(filterLayout)

        self.chipNumberLine = QtGui.QLineEdit(str(self.chip))
        self.chipNumberLine.setMaxLength(2)
        self.decChipButton = QtGui.QPushButton()
        left_arrow_path = os.path.join(os.path.dirname(__file__),
                                        '../img/left_arrow.png')
        self.decChipButton.setIcon(QtGui.QIcon(left_arrow_path))
        self.decChipButton.clicked.connect(self.decrementChip)
        self.setChipButton = QtGui.QPushButton('Plot chip specified')
        self.setChipButton.clicked.connect(self.changeChip)
        self.chipSelection = QtGui.QWidget()
        self.incChipButton = QtGui.QPushButton()
        right_arrow_path = os.path.join(os.path.dirname(__file__),
                                        '../img/right_arrow.png')
        self.incChipButton.setIcon(QtGui.QIcon(right_arrow_path))
        self.incChipButton.clicked.connect(self.incrementChip)
        chipButtonsLayout = QtGui.QHBoxLayout()
        chipButtonsLayout.addWidget(self.decChipButton)
        chipButtonsLayout.addWidget(self.setChipButton)
        chipButtonsLayout.addWidget(self.incChipButton)
        chipLayout = QtGui.QVBoxLayout()
        chipLayout.addWidget(self.chipNumberLine)
        chipLayout.addLayout(chipButtonsLayout)
        self.chipSelection.setLayout(chipLayout)

        if self.impedance != None:
            self.impedanceCheckbox = QtGui.QCheckBox('Display impedance')
            self.impedanceCheckbox.setCheckState(QtCore.Qt.Checked)
            self.impedanceCheckbox.toggled.connect(self.toggleImpedanceLabels)

            self.refCheckbox = QtGui.QCheckBox('Use software reference')
            self.refCheckbox.setCheckState(QtCore.Qt.Unchecked)
            self.refCheckbox.toggled.connect(self.toggleRef)
        self.use_ref = False

        self.buttonPanel = QtGui.QWidget()
        buttonsLayout = QtGui.QHBoxLayout()
        buttonsLayout.addWidget(self.startButton)
        buttonsLayout.addWidget(self.stopButton)
        buttonsLayout.addWidget(self.chipSelection)
        buttonsLayout.addWidget(self.filterPanel)
        if self.impedance != None:
            buttonsLayout.addWidget(self.refCheckbox)
            buttonsLayout.addWidget(self.impedanceCheckbox)
        self.buttonPanel.setLayout(buttonsLayout)
        self.buttonPanel.setMaximumHeight(150)

        maxWidth = 800
        if self.impedance != None:
            maxWidth += 400
        self.buttonPanel.setMaximumWidth(maxWidth)

        self.layout = QtGui.QVBoxLayout()
        self.layout.addWidget(self.buttonPanel)
        self.layout.addWidget(self.labels)
        self.layout.addWidget(self.canvas_embedded.native)

        self.setLayout(self.layout)

        self.setWindowTitle('Willow Live Streaming')
        self.setWindowIcon(QtGui.QIcon('../img/round_logo_60x60.png'))

        self.canvas_popout.native.hide()

        self.changeChipTimer = QtCore.QTimer()
        self.changeChipTimer.timeout.connect(self.changeChip)
        self.changeChipTimer.setSingleShot(True)

    def canvasesApply(self, function, *args, **kwargs):
        embedded_args = (self.canvas_embedded,) + args
        function(*embedded_args, **kwargs)
        # avoid updating canvas_popout widget while hidden to bound memory use
        if (self.canvas_popout.native.isVisible() or
            function != ChipCanvas.draw_new_data):
            popout_args = (self.canvas_popout,) + args
            function(*popout_args, **kwargs)

    def toggleFiltering(self, checkbox_state):
        self.filtering = checkbox_state

    def toggleRef(self, checkbox_state):
        self.use_ref = checkbox_state

    def toggleImpedanceLabels(self, checkbox_state):
        self.display_impedance = checkbox_state
        self.canvasesApply(ChipCanvas.updateLabels)
        self.canvasesApply(ChipCanvas.update)

    def setLowcut(self):
        self.lowcut = float(self.lowcutLine.text())

    def setHighcut(self):
        self.highcut = float(self.highcutLine.text())

    def launchFilterPlot(self):
        filt = self.getFilter()
        self.fpw = FilterPlotWindow(filt, [self.lowcut, self.highcut])
        self.fpw.show()

    def hwif_request(self, *args):
        """Write to stdout requests to Willow hardware; WillowGUI will execute
        these requests if it can. """
        sys.stderr.write('hwif_req: ' + ', '.join(args) + '\n')

    def startStreaming(self):
        self.hwif_request('startStreaming_subsamples')
        self.streaming = True
        print "spawning proto2bytes subprocess"
        self.p2b_po = subprocess.Popen([self.p2b_filename, '-A'],
            stdout=subprocess.PIPE,
            # comment this out to view stderr in terminal (kinda useful for debugging)
            stderr=subprocess.PIPE
            )
        # set proto2bytes' stdout as non-blocking
        fl = fcntl.fcntl(self.p2b_po.stdout.fileno(), fcntl.F_GETFL)
        fcntl.fcntl(self.p2b_po.stdout.fileno(), fcntl.F_SETFL, fl | os.O_NONBLOCK)
        # poller for telling if reads are ready (POLLIN), proto2bytes has
        # stopped (POLLHUP), or if it is having an error (POLLERR)
        self.p2b_poller = select.poll()
        self.p2b_poller.register(self.p2b_po.stdout,
                select.POLLIN | select.POLLERR | select.POLLHUP)
        # FileIO created from stdout (used for readinto() calls when poling)
        self.p2b_stdoutio = io.FileIO(self.p2b_po.stdout.fileno(), mode='r')
        self.pollStream_bytes_so_far = 0
        self.pollStream_failed_polls = 0
        self.pollStreamTimer = QtCore.QTimer()
        self.pollStreamTimer.timeout.connect(self.pollStream)
        # poll whenever rest of QT events are not being serviced
        self.pollStreamTimer.start(0)

    def stopStreaming(self):
        self.hwif_request('stopStreaming')
        self.streaming = False
        try:
            self.p2b_po.kill()
            self.pollStreamTimer.stop()
        except AttributeError:
            pass

    def playPause(self):
        if self.streaming:
            self.stopStreaming()
        else:
            self.startStreaming()

    def pollStream(self):
        if not self.streaming:
            return
        if self.pollStream_bytes_so_far != len(self.new_data_ba):
            pollout = self.p2b_poller.poll(POLL_TIME_MS)
            if pollout:
                self.pollStream_failed_polls = 0
                [(_, events)] = pollout
                # if proto2bytes failing to stream, give up
                if select.POLLIN & events == select.POLLIN:
                    rest_of_buff = self.new_data_mv[self.pollStream_bytes_so_far:]
                    self.pollStream_bytes_so_far += self.p2b_stdoutio.readinto(rest_of_buff)
                if select.POLLERR & events == select.POLLERR:
                    sys.stderr.write("Chipstream's proto2bytes in error state!\n")
                    self.stopStreaming()
                if select.POLLHUP & events == select.POLLHUP:
                    sys.stderr.write("Chipstream's proto2bytes hung up!\n")
                    self.stopStreaming()
                if self.pollStream_bytes_so_far == len(self.new_data_ba):
                    # whole frame of data ready! let's draw it!
                    self.updatePlot()
                    self.pollStream_bytes_so_far = 0
            else:
                self.pollStream_failed_polls += 1
                if self.pollStream_failed_polls * POLL_TIME_MS >= STREAM_GIVEUP_TIME_MS:
                    sys.stderr.write("Chipstream's proto2bytes failed to give data for %d seconds! Killing it...\n" \
                                     % STREAM_GIVEUP_TIME_MS / 1000.)
                    self.stopStreaming()

    def setInfoLabel(self, xr, yr):
        self.infoLabel.setText('Chip %d: xrange = %3.2f ms, yrange = %3.2f uV\n'
                               'mousewheel to zoom x and y, ctrl+mousewheel to zoom x, shift+mousewheel to zoom y'
                                        % (self.chip, xr, yr))
    def setPlotLabel(self, chan):
        if self.probeMap != None:
            if chan != None:
                for key, value in self.probeMap.items():
                    if value == chan and isinstance(key, tuple):
                        (shank, row, col) = key
                self.plotLabel.setText('Channel %d selected\n'
                    '(shank %d, row %d, column %d)' % (chan, shank, row, col))
            else:
                self.plotLabel.setText('')

    def getFilter(self):
        nyq = 1.5e4    # 30 kHz sampling frequency -> 15 kHz nyquist rate
        order = 5
        warnings.filterwarnings('error') # catch division by zero when filtering
        try:
            b, a = signal.butter(order,[self.lowcut/nyq, self.highcut/nyq],btype='bandpass')
            warnings.filterwarnings('default')
            return b, a
        except:
            raise Exception('Illegal filtering cutoffs (low={}, high={}).'.format(
                   self.lowcut, self.highcut))
            warnings.filterwarnings('default')
            return 1, 1

    def incrementChip(self):
        chip = min (int(self.chipNumberLine.text()) + 1, MAX_CHIP_NUMBER)
        self.chipNumberLine.setText(str(chip))
        # if not currently aggregating changes, instantiate immediately
        if not self.changeChipTimer.isActive():
            self.changeChip()
        # arm or rearm the aggregation timer
        self.changeChipTimer.start(CHANGE_CHIP_MS)

    def decrementChip(self):
        chip = max(int(self.chipNumberLine.text()) - 1, MIN_CHIP_NUMBER)
        self.chipNumberLine.setText(str(chip))
        # if not currently aggregating changes, instantiate immediately
        if not self.changeChipTimer.isActive():
            self.changeChip()
        # arm or rearm the aggregation timer
        self.changeChipTimer.start(CHANGE_CHIP_MS)

    def changeChip(self):
        chip = int(self.chipNumberLine.text())
        if chip != self.chip:
            self.hwif_request('setSubsamples_byChip', str(chip))
            self.draw_point = 0
            self.initChip(chip)

    def initChip(self, chip):
        self.chip = chip
        if self.impedance != None:
            self.bad_chans = np.array([c for c in range(nchans) if
                self.impedanceBad(self.getImpedance(c+self.chip*nchans))])
            self.impMask = np.zeros_like(self.plotBuff_raw)
            self.impMask[self.bad_chans,:] = True
        self.canvasesApply(ChipCanvas.updateLabels)
        self.canvasesApply(ChipCanvas.update)
        self.canvasesApply(ChipCanvas.draw_new_data, np.zeros((nchans, self.nsamps), dtype='float'),
                        self.nsamps)

        self.setInfoLabel(self.deltax_ms, self.deltay_uv)

    def getImpedance(self, chan_idx):
        return self.impedance[chan_idx]/1000.   # in kohms

    def impedanceBad(self, impedance):
        # tests whether an impedance, in kohms, is too high
        # (indicating open-circuit channel)
        return impedance > 1000. or impedance < 100.


    def updatePlot(self):
        try:
            # assumes daemon running on a little-endian system
            # (true of x86, amd64, but not of some other processors)
            new_data = struct.unpack('<'+'H'*nchans*self.nrefresh, self.new_data_ba)
        except (struct.error, TypeError) as e:
            sys.stderr.write('\n' + str(e) + '\n')
            sys.stderr.write('Read from proto2bytes failed!\n')
            self.stopStreaming()
            return

        for i in xrange(self.nrefresh):
            self.newBuff[:,i] = (np.array(new_data[i*nchans:(i+1)*nchans], dtype='float')-2**15)*MICROVOLTS_PER_COUNT

        if self.draw_point == 0:
            self.plotBuff_raw = np.concatenate((self.newBuff[:,:],
                self.plotBuff_raw[:,self.nrefresh:]),1)
            last_point = self.draw_point + self.nrefresh
        elif self.draw_point + self.nrefresh > self.nsamps:
            # overflow of new buffer back to beginning of plot
            overflow_amount = self.draw_point + self.nrefresh - self.nsamps
            self.plotBuff_raw = np.concatenate((self.newBuff[:,-overflow_amount:],
                self.plotBuff_raw[:,overflow_amount:self.draw_point],
                self.newBuff[:,:self.nrefresh-overflow_amount]),1)
            last_point = overflow_amt
        else:
            self.plotBuff_raw = np.concatenate((self.plotBuff_raw[:,:self.draw_point],
                self.newBuff[:,:],
                self.plotBuff_raw[:,self.draw_point+self.nrefresh:]),1)
            last_point = self.draw_point + self.nrefresh

        ########
        # post-processing
        #####

        if self.use_ref:
            m_plotBuff = np.ma.masked_array(self.plotBuff_raw, self.impMask)
            mean = np.ma.mean(m_plotBuff, axis=0)
            if np.ma.is_masked(mean):
                # if all channels to display were bad
                self.plotBuff = self.plotBuff_raw
            else:
                # subtract mean data from good channels (leave bad channels)
                good_chans = [c for c in range(nchans) if c not in self.bad_chans]
                self.plotBuff[good_chans,:] = self.plotBuff_raw[good_chans,:] - mean
                self.plotBuff[self.bad_chans,:] = self.plotBuff_raw[self.bad_chans,:]
        else:
            self.plotBuff = self.plotBuff_raw

        if self.filtering:
            b, a = self.getFilter()
            self.plotBuff = signal.filtfilt(b, a, self.plotBuff, axis=1)

        self.canvasesApply(ChipCanvas.draw_new_data, self.plotBuff, last_point)
        self.draw_point += self.nrefresh
        self.draw_point %= self.plotBuff.shape[1]

    def popoutCanvas(self):
        self.canvas_popout.native.showFullScreen()
        self.poppedIn = False

    def popinCanvas(self):
        self.canvas_popout.native.hide()
        self.poppedIn = True

    def keyReleaseEvent(self, event):
        if not self.poppedIn and (event.key == QtCore.Qt.Key_F11 or
            event.key == QtCore.Qt.Key_Escape):
            self.popinCanvas()
        elif self.poppedIn and event.key == QtCore.Qt.Key_F11:
            self.popoutCanvas()
        elif event.key == QtCore.Qt.Key_Space:
            self.playPause()


    def onClose(self, event):
        self.canvas_popout.native.hide()
        self.canvas_popout.native.close()
        self.canvas_popout.close()

VERT_SHADER = """
#version 120

// y coordinate of the position.
attribute float a_position;

// Range of displayable y-coordinates.
uniform vec2 u_yrange;

// Row, col, and time index.
attribute vec3 a_index;
varying vec3 v_index;

// 2D scaling factor (zooming).
uniform vec2 u_scale;

// Number of samples per signal.
uniform float u_nsamps;

// Color.
uniform vec4 u_colors[{NCHANNELS}];
varying vec4 v_color;

// Index in plot representing where new data begins to be drawn.
// (Used for coloring eye-candy.)
uniform int u_latest_point;

// Varying variables used for clipping in the fragment shader.
varying vec2 v_position;
varying vec4 v_ab;

void main() {{
    float ymin = u_yrange.x;
    float ymax = u_yrange.y;
    // Normalize data, with x in the range (0, u_nsamps-1), y in range (ymin, ymax),
    // to fit within a subplot, with x in range (-1,1), y in range(-1,1).
    float x = -1 + 2*a_index.z / (u_nsamps-1);
    float y = 0 + 2*a_position / (ymax-ymin);
    vec2 position = vec2(x, y);
    // Find the affine transformation for the subplots.
    vec2 a = vec2(1./{NCOLS}, 1./{NROWS})*{SUBPLOT_HORIZONTAL_FILL};
    vec2 b = vec2(-1 + 2*(a_index.x + 0.5) / {NCOLS},
                   1 - 2*(a_index.y + 0.5) / {NROWS});
    // Apply the static subplot transformation + scaling.
    gl_Position = vec4(a*u_scale*position+b, 0.0, 1.0);

    int channel = int(a_index.y*{NCOLS} + a_index.x);

    v_color = u_colors[channel];
    // Make most recent data brighter and least-recent data dimmer.
    // This makes following new data easier on eyes.
    float oldness;
    if ((a_index.z <= u_latest_point))
        oldness = u_latest_point - a_index.z;
    else
        oldness = u_latest_point + (u_nsamps - a_index.z);
    v_color.x = v_color.x * 1.0 - (oldness/u_nsamps)*0.60;
    v_color.y = v_color.y * 1.0 - (oldness/u_nsamps)*0.60;
    v_color.z = v_color.z * 1.0 - (oldness/u_nsamps)*0.60;
    v_index = a_index;
    // For clipping test in the fragment shader.
    v_position = gl_Position.xy;
    v_ab = vec4(a, b);
}}
""".format(NCOLS=ncols, NROWS=nrows, NCHANNELS=nrows*ncols,
        SUBPLOT_HORIZONTAL_FILL=0.95)

FRAG_SHADER = """
#version 120

varying vec4 v_color;
varying vec3 v_index;

varying vec2 v_position;
varying vec4 v_ab;

void main() {
    gl_FragColor = v_color;
    // Discard the fragments between the signals (emulate glMultiDrawArrays).
    if ((fract(v_index.x) > 0.) || (fract(v_index.y) > 0.))
        discard;
    // Clipping test.
    vec2 test = abs((v_position.xy-v_ab.zw)/v_ab.xy);
    if ((test.x > 1) || (test.y > 1))
        discard;
}
"""

class ChipCanvas(app.Canvas):

    def __init__(self, nsamps, nrefresh, ymin, ymax, ctl_panel, parent=None):
        super(ChipCanvas, self).__init__(app='pyqt4', keys=None, parent=parent)

        self.ctl_panel = ctl_panel
        self.parent = parent        # if is None, canvas in own window. if is ctl_panel, canvas embedded.
        self.nrefresh = nrefresh
        self.nsamps = nsamps
        self.deltay_uv = ymax - ymin

        vertex_data = np.zeros((ncols, nrows, self.nsamps),
            dtype=[('a_index', np.float32, 3)])

        self.subplot_indices = set((x,y) for x,y in \
            itertools.product(xrange(ncols),xrange(nrows)))

        point_indices = set((x,y,i) for x,y,i in \
            itertools.product(xrange(ncols),xrange(nrows),xrange(self.nsamps)))
        for (x,y,i) in point_indices:
            vertex_data['a_index'][x,y,i] = np.array([x,y,i])

        self._program = gloo.Program(VERT_SHADER, FRAG_SHADER)
        self._program.bind(gloo.VertexBuffer(vertex_data))

        # a_position field will be yvals of subplots, in uVolts
        self.repacked_yvals = np.zeros((ncols, nrows, self.nsamps), np.float32)

        self._program['a_position'] = self.repacked_yvals.ravel()

        self._program['u_latest_point'] = self.nrefresh
        self._program['u_yrange'] = (ymin, ymax)
        self.scale_x = 1.
        self.scale_y = 1.
        self._program['u_scale'] = (self.scale_x, self.scale_y)

        self.real_scale = (int(10 * round(float(self.nsamps/(30*1.))/10)),
                           int(10 * round(float(self.deltay_uv/1.)/10)))

        self._program['u_nsamps'] = self.nsamps

        self.default_rgb = np.array((0.559, 0.855, 0.563))
        self.bad_impedance_rgb = np.array((0.5, 0.5, 0.5))
        self.subplot_indices = set((x,y) for x,y in \
            itertools.product(xrange(ncols),xrange(nrows)))

        default_rgba = np.append(self.default_rgb, 1.0)
        for (x,y) in self.subplot_indices:
            chan_idx = y*ncols + x
            self._program['u_colors[{CI}]'.format(CI=chan_idx)] = default_rgba

        # borders around plots
        self.border_collection = None

        # labels around plots
        self.label_visual = None

        # axis to be drawn around plot mouse clicks on
        self.selected_plots = []
        self.axis_plot = None
        self.axis_scale = self.real_scale
        self.axis_y = None
        self.show_axis = False

        gloo.set_viewport(0, 0, *self.physical_size)

        gloo.set_state(clear_color=(0.10, 0.10, 0.10, 1.0), blend=True,
                       blend_func=('src_alpha', 'one_minus_src_alpha'))

        self.createBorders()

    def hasParent(self):
        return (self.parent is not None)

    def createBorders(self):
        plot_w = 2.0/ncols
        plot_h = 2.0/nrows

        n_xborders = ncols - 1
        n_yborders = nrows - 1

        self.border_collection = SegmentCollection("agg")

        # horizontal borders
        c0 = np.dstack(
             (np.linspace(-1.0+plot_w, 1.0-plot_w, n_xborders), \
              -1.0*np.ones(n_xborders), np.zeros(n_xborders))).reshape(n_xborders, 3)
        c1 = np.dstack(
             (np.linspace(-1.0+plot_w, 1.0-plot_w, n_xborders), \
              np.ones(n_xborders), np.zeros(n_xborders))).reshape(n_xborders, 3)
        self.border_collection.append(c0, c1)

        # vertical borders
        r0 = np.dstack(
             (-1.0*np.ones(n_yborders), \
              np.linspace(-1.0+plot_h, 1.0-plot_h, n_yborders), \
              np.zeros(n_yborders))).reshape(n_yborders, 3)
        r1 = np.dstack(
             (np.ones(n_yborders), \
              np.linspace(-1.0+plot_h, 1.0-plot_h, n_yborders), \
              np.zeros(n_yborders))).reshape(n_yborders, 3)
        self.border_collection.append(r0, r1)

        self.border_collection['antialias'] = 1
        self.border_collection['viewport'] = 0, 0, self.physical_size[0], self.physical_size[1]
        self.border_collection['color'] = (1.0, 1.0, 1.0, 1.0)

    def reportLocation(self, plot):
        if plot != None:
            plot_x, plot_y = plot
            chan_idx = plot_y*ncols + plot_x + nchans*self.ctl_panel.chip
            self.ctl_panel.setPlotLabel(chan_idx)
        else:
            self.ctl_panel.setPlotLabel(None)

    def createLabels(self):
        self.label_visual = visuals.TextVisual(text='',
            color=(0.5,0.5,0.5,1.0), font_size=8, pos=(0,0), anchor_x='left')
        self.label_visual.transform = transforms.NullTransform()
        self.updateLabels()

    def updateLabels(self):
        if self.label_visual is not None:
            texts = []
            poses = []
            a = (self.physical_size[0]/ncols, self.physical_size[1]/nrows)
            for (x,y) in self.subplot_indices:
                chan_idx = y*ncols + x + nchans*self.ctl_panel.chip
                texts.append(' Channel {c}'.format(c=chan_idx))
                poses.append((a[0]*x, a[1]*(y+0.125)))
                if self.ctl_panel.display_impedance:
                    impedance = self.ctl_panel.getImpedance(chan_idx)
                    impedance_label = u' Impedance: {i:,.2f} k\N{OHM SIGN}'.format(i=impedance)
                    if self.ctl_panel.impedanceBad(impedance):
                        impedance_label += ' (BAD)'
                    texts.append(impedance_label)
                    poses.append((a[0]*x, a[1]*(y+0.875)))
            self.label_visual.text = texts
            self.label_visual.pos  = poses
        else:
            self.createLabels()

    def clearAxis(self):
        self.axis_plot = None
        self.axis_scale = None
        self.axis_y = None

    def createAxis(self, plot):
        if self.axis_plot == plot and self.real_scale == self.axis_scale:
            print 'redundant'
            return
        print 'making new axis'

        plot_w = 1.0/ncols
        plot_h = 1.0/nrows
        x_margin = 0.025 # plots have 0.025*plot_w margins to the left & right
        plot_x, plot_y = plot
        pos_xax = np.array([[plot_w*(plot_x + x_margin),
                             plot_h*(plot_y + 0.5)],
                            [plot_w*(plot_x + 1.0 - x_margin),
                             plot_h*(plot_y + 0.5)]])
        pos_xax *= np.array([self.physical_size[0], self.physical_size[1]])

        #time_visible  = self.real_scale[0]    # in ms
        #time_0 , time_1 = 0.0, time_visible
        #self.axis_x = visuals.AxisVisual(pos_xax,
        #    domain=(time_0, time_1), tick_direction=(0., 1.),
        #    axis_color=(0.5,0.5,0.5), tick_color=(0.4,0.4,0.4))

        pos_yax = np.array([[plot_w*(plot_x+1-x_margin),
                             plot_h*(plot_y)],
                            [plot_w*(plot_x+1-x_margin),
                             plot_h*(plot_y + 1.0)]])
        pos_yax *= np.array([self.physical_size[0], self.physical_size[1]])

        uVolts_0, uVolts_1 = (self.real_scale[1], -self.real_scale[1])
        self.axis_y = visuals.AxisVisual(pos_yax,
            domain=(uVolts_0, uVolts_1), tick_direction=(-1., 0.),
            axis_color=(0.5,0.5,0.5), tick_color=(0.4,0.4,0.4))

        vp = 0, 0, self.physical_size[0], self.physical_size[1]
        #self.axis_x.transforms.configure(canvas=self, viewport=vp)
        self.axis_y.transforms.configure(canvas=self, viewport=vp)
        self.axis_plot = plot
        self.axis_scale = self.real_scale

    def focusPlots(self):
        if self.selected_plots != []:
            # unfocus non-selected subplots
            unfocused_color = np.append(((0.75/0.855)*self.default_rgb), 1.0)
            bad_impedance_color = np.append(self.bad_impedance_rgb, 1.0)
            for (x,y) in self.subplot_indices:
                chan_idx = y*ncols + x + nchans*self.ctl_panel.chip
                plot_idx = y*ncols + x
                if self.ctl_panel.display_impedance and self.ctl_panel.impedanceBad(self.ctl_panel.getImpedance(chan_idx)):
                    self._program['u_colors[{C}]'.format(C=plot_idx)] = bad_impedance_color
                else:
                    self._program['u_colors[{C}]'.format(C=plot_idx)] = unfocused_color

            # focus selcted plots
            for plot in self.selected_plots:
                plot_x, plot_y = plot
                selected_idx = plot_y*ncols + plot_x
                focused_color = (1.0, 1.0, 1.0, 1.0)
                self._program['u_colors[{C}]'.format(C=selected_idx)] = focused_color
        else:
            default_color = np.append(self.default_rgb, 1.0)
            bad_impedance_color = np.append(self.bad_impedance_rgb, 1.0)
            for (x,y) in self.subplot_indices:
                chan_idx = y*ncols + x + nchans*self.ctl_panel.chip
                plot_idx = y*ncols + x
                if self.ctl_panel.display_impedance and self.ctl_panel.impedanceBad(self.ctl_panel.getImpedance(chan_idx)):
                    self._program['u_colors[{C}]'.format(C=plot_idx)] = bad_impedance_color
                else:
                    self._program['u_colors[{C}]'.format(C=plot_idx)] = default_color

    # used by parent widget. new_yvals comes in as an array with shape (nchans, nsamps)
    def draw_new_data(self, new_yvals, latest_point):
        for chan in xrange(ncols*nrows):
            self.repacked_yvals[chan%ncols, chan//ncols,:] = new_yvals[chan,:]
        self._program['a_position'] = self.repacked_yvals.ravel().astype(np.float32)
        self._program['u_latest_point'] = latest_point
        self.focusPlots()
        self.update()

    def on_resize(self, event):
        vp = 0, 0, self.physical_size[0], self.physical_size[1]
        self.context.set_viewport(*vp)
        self.border_collection['viewport'] = vp
        self.updateLabels()
        self.label_visual.transforms.configure(canvas=self, viewport=vp)
        self.update()

    def on_draw(self, event):
        gloo.clear()
        self._program.draw('line_strip')
        self.border_collection.draw()
        if self.show_axis and self.selected_plots != []:
            self.createAxis(self.selected_plots[-1])
            self.axis_y.draw()
        self.label_visual.draw()

    def on_mouse_wheel(self, event):
        d = np.sign(event.delta[1]) * .05
        if 'Shift' in event.modifiers:
            self.scale_x, self.scale_y = (max(1, self.scale_x*np.exp(0.0*d)),
                                          max(1, self.scale_y*np.exp(2.5*d)))
        elif 'Control' in event.modifiers:
            self.scale_x, self.scale_y = (max(1, self.scale_x*np.exp(2.5*d)),
                                          max(1, self.scale_y*np.exp(0.0*d)))
        else:
            self.scale_x, self.scale_y = (max(1, self.scale_x*np.exp(2.5*d)),
                                          max(1, self.scale_y*np.exp(2.5*d)))

        # axes drawn too many times if we are scrolling more than one teeny bit
        # at a time, slowing down axis scaling when zooming. for now, just
        # remove axes after scrolling and re-enable them to see new measurement
        self.show_axis = False
        self.clearAxis()

        # constrain scale factors
        range_x_new = self.nsamps/(30*self.scale_x)
        if range_x_new < self.ctl_panel.frame_period:
            range_x_new = self.ctl_panel.frame_period
            self.scale_x = self.nsamps/(30*range_x_new)
        range_y_new = self.deltay_uv/self.scale_y
        if range_y_new < MICROVOLTS_PER_COUNT:
            range_y_new = MICROVOLTS_PER_COUNT
            self.scale_y = self.deltay_uv/range_y_new

        self._program['u_scale'] = (self.scale_x, self.scale_y)

        # round new scales to the nearest 10 uV for cleanliness
        scale_x_real = int(10 * round(float(self.nsamps/(30*self.scale_x))/10))
        scale_y_real = int(10 * round(float(self.deltay_uv/self.scale_y)/10))

        self.real_scale = (self.scale_x, self.scale_y)

        self.ctl_panel.setInfoLabel(scale_x_real, scale_y_real)

        self.update()

    def on_mouse_press(self, click):
        # plot that mouse is within
        plot_x = ncols*click.pos[0]/self.size[0]
        plot_y = nrows*click.pos[1]/self.size[1]
        plot = (plot_x, plot_y)

        modifiers = QtGui.QApplication.keyboardModifiers()

        if modifiers & QtCore.Qt.ControlModifier:
            # focus multiple subplots when Ctrl held down
            if plot not in self.selected_plots:
                self.selected_plots.append(plot)
                self.reportLocation(plot)
            else:
                self.reportLocation(None)
                self.selected_plots.remove(plot)
        else:
            # focus only plot most recently clicked when Ctrl not held down
            if plot not in self.selected_plots:
                self.selected_plots = [plot]
                self.reportLocation(plot)
            else:
                self.selected_plots = []
                self.reportLocation(None)
        self.focusPlots()
        self.update()

    def on_key_press(self, event):
        # 'a' key toggles axis display on most recently selected plot
        if event.key.name == 'A' and self.selected_plots != []:
            self.show_axis = not self.show_axis
            if self.show_axis == False:
                #self.clearAxis()
                self.axis_plot = None
                self.axis_scale = None
                self.axis_y = None
                self.update()
            else:
                self.createAxis(self.selected_plots[-1])
                self.update()
        elif (not self.hasParent()) and (event.key.name == 'Escape' or event.key.name == 'F11'):
            self.ctl_panel.popinCanvas()
        elif (self.hasParent()) and event.key.name == 'F11':
            self.ctl_panel.popoutCanvas()
        elif event.key.name == 'Escape':
            pass    # if we didn't do this, Esc would close the embedded canvas
        elif event.key.name == 'Space':
            self.ctl_panel.playPause()

if __name__=='__main__':
    from StreamDialog import StreamDialog
    if len(sys.argv) < 2:
        print "Need location of proto2bytes executable as first command-line argument."
        sys.exit()
    proto2bytes_path = str(sys.argv[1])
    app = QtGui.QApplication(sys.argv)
    dlg = StreamDialog()
    if dlg.exec_():
        params = dlg.getParams()
        streamWindow = StreamWindow(params, proto2bytes_path)
        streamWindow.show()
    sys.exit(app.exec_())
