#!/usr/bin/env python

try:
    import scipy
    from scipy import fftpack
except ImportError:
    print "Please install SciPy to run this script (http://www.scipy.org/)"
    raise SystemExit, 1

import sys, os
from PyQt4 import Qt, QtCore, QtGui
import PyQt4.Qwt5 as Qwt
from matplotlib import mlab
from optparse import OptionParser
from gnuradio import eng_notation

from pyqt_plot import Ui_MainWindow

class gr_plot_qt(QtGui.QMainWindow):
    def __init__(self, qapp, filename, options, parent=None):
        QtGui.QWidget.__init__(self, parent)
        self.gui = Ui_MainWindow()
        self.gui.setupUi(self)
                       
        self.block_length = options.block_length
        self.start = options.start
        self.sample_rate = options.sample_rate
        self.psdfftsize = options.psd_size
        self.specfftsize = options.spec_size
        self.winfunc = scipy.blackman
        self.sizeof_data = 8
        self.datatype = scipy.complex64
        self.iq = list()
        self.time = list()
        
        # Set up basic plot attributes
        self.gui.timePlot.setAxisTitle(self.gui.timePlot.xBottom, "Time (sec)")
        self.gui.timePlot.setAxisTitle(self.gui.timePlot.yLeft, "Amplitude (V)")
        self.gui.freqPlot.setAxisTitle(self.gui.freqPlot.xBottom, "Frequency (Hz)")
        self.gui.freqPlot.setAxisTitle(self.gui.freqPlot.yLeft, "Magnitude (dB)")

        # Set up FFT size combo box
        self.gui.fftComboBox.addItems(["128", "256", "512", "1024", "2048",
                                       "4096", "8192", "16384", "32768"])
        pos = self.gui.fftComboBox.findText(Qt.QString("%1").arg(self.psdfftsize))
        self.gui.fftComboBox.setCurrentIndex(pos)
        self.connect(self.gui.fftComboBox,
                     Qt.SIGNAL("activated (const QString&)"),
                     self.fftComboBoxEdit)

        # Set up color scheme box
        self.color_modes = {"Black on White" : self.color_black_on_white,
                            "White on Black" : self.color_white_on_black,
                            "Blue on Black"  : self.color_blue_on_black,
                            "Green on Black" : self.color_green_on_black}
        self.gui.colorComboBox.addItems(self.color_modes.keys())
        pos = self.gui.colorComboBox.findText("Blue on Black")
        self.gui.colorComboBox.setCurrentIndex(pos)
        self.connect(self.gui.colorComboBox,
                     Qt.SIGNAL("activated (const QString&)"),
                     self.colorComboBoxEdit)
        
        
        # Create zoom functionality for the plots
        self.timeZoomer = Qwt.QwtPlotZoomer(self.gui.timePlot.xBottom,
                                            self.gui.timePlot.yLeft,
                                            Qwt.QwtPicker.PointSelection,
                                            Qwt.QwtPicker.AlwaysOn,
                                            self.gui.timePlot.canvas())

        self.freqZoomer = Qwt.QwtPlotZoomer(self.gui.freqPlot.xBottom,
                                            self.gui.freqPlot.yLeft,
                                            Qwt.QwtPicker.PointSelection,
                                            Qwt.QwtPicker.AlwaysOn,
                                            self.gui.freqPlot.canvas())

        self.picker = Qwt.QwtPlotPicker(self.gui.timePlot.xBottom,
                                        self.gui.timePlot.yLeft,
                                        Qwt.QwtPicker.PointSelection,
                                        Qwt.QwtPlotPicker.CrossRubberBand,
                                        Qwt.QwtPicker.AlwaysOn,
                                        self.gui.timePlot.canvas())
        self.picker.connect(self.picker,
                            Qt.SIGNAL('selected(const QwtDoublePoint&)'),
                            self.clickMe)

        # Set up action when tab is changed
        self.connect(self.gui.tabGroup,
                     Qt.SIGNAL("currentChanged (int)"),
                     self.tabChanged)

        # Add a legend to the Time plot
        legend_real = Qwt.QwtLegend()
        self.gui.timePlot.insertLegend(legend_real)

        # Set up slider
        self.gui.plotHBar.setSingleStep(1)
        self.gui.plotHBar.setPageStep(self.block_length)
        self.gui.plotHBar.setMinimum(0)
        self.gui.plotHBar.setMaximum(self.block_length)
        self.connect(self.gui.plotHBar,
                     Qt.SIGNAL("valueChanged(int)"),
                     self.sliderMoved)

        # Connect Open action to Open Dialog box
        self.connect(self.gui.action_open,
                     Qt.SIGNAL("activated()"),
                     self.open_file)
        
        # Set up file position boxes to update current figure
        self.connect(self.gui.filePosStartLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_position_changed)
        self.connect(self.gui.filePosStopLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_position_changed)
        self.connect(self.gui.filePosLengthLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_length_changed)

        self.connect(self.gui.fileTimeStartLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_time_changed)
        self.connect(self.gui.fileTimeStopLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_time_changed)
        self.connect(self.gui.fileTimeLengthLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.file_time_length_changed)

        self.rcurve = Qwt.QwtPlotCurve("Real")
        self.icurve = Qwt.QwtPlotCurve("Imaginary")

        self.icurve.attach(self.gui.timePlot)
        self.rcurve.attach(self.gui.timePlot)

        self.psdcurve = Qwt.QwtPlotCurve("PSD")
        self.psdcurve.attach(self.gui.freqPlot)

        # Set up initial color scheme
        self.color_modes["Blue on Black"]()

        self.set_sample_rate(self.sample_rate)
        self.connect(self.gui.sampleRateLineEdit,
                     Qt.SIGNAL("editingFinished()"),
                     self.sample_rate_changed)

        if(filename is not None):
            self.initialize(filename)

        self.show()

    def open_file(self):
        filename = Qt.QFileDialog.getOpenFileName(self, "Open", ".")
        print filename
        self.initialize(filename)

    def initialize(self, filename):
        self.hfile = open(filename, "r")

        self.setWindowTitle(("GNU Radio File Plot Utility: %s" % filename))
        
        self.cur_start = 0
        self.cur_stop = self.block_length

        self.init_data_input()
        self.get_data(self.cur_start, self.cur_stop)
        self.get_psd()
        self.gui.plotHBar.setSliderPosition(0)
        self.gui.plotHBar.setMaximum(self.signal_size)

        self.update_time_curves()
        self.update_psd_curves()

    def init_data_input(self):
        self.hfile.seek(0, os.SEEK_END)
        self.signal_size = self.hfile.tell()/self.sizeof_data
        print "Sizeof File: ", self.signal_size
        self.hfile.seek(0, os.SEEK_SET)
        
    def get_data(self, start, end):
        if(end > start):
            self.hfile.seek(start*self.sizeof_data, os.SEEK_SET)
            self.position = start
            iq = scipy.fromfile(self.hfile, dtype=self.datatype,
                                count=end-start)
            if(len(iq) < (end-start)):
                print "End of File"
            else:
                tstep = 1.0 / self.sample_rate
                self.iq = iq
                self.time = [tstep*(self.position + i) for i in xrange(len(self.iq))]

            self.set_file_pos_box(start, end)
        else:
            # Do we want to do anything about this?
            pass

    def get_psd(self):
        winpoints = self.winfunc(self.psdfftsize)
        iq_psd, freq = mlab.psd(self.iq, Fs=self.sample_rate,
                                NFFT=self.psdfftsize,
                                noverlap=self.psdfftsize/4.0,
                                window=winpoints,
                                scale_by_freq=False)
        self.iq_psd = 10.0*scipy.log10(abs(fftpack.fftshift(iq_psd)))
        self.freq = freq - self.sample_rate/2.0


    def clickMe(self, qPoint):
        print qPoint.x()

    def fftComboBoxEdit(self, fftSize):
        self.psdfftsize = fftSize.toInt()[0]
        self.get_psd()
        self.update_psd_curves()
        
    def colorComboBoxEdit(self, colorSelection):
        colorstr = str(colorSelection.toAscii())
        color_func = self.color_modes[colorstr]
        color_func()

    def sliderMoved(self, value):
        pos_start = value
        pos_end = value + self.gui.plotHBar.pageStep()

        self.get_data(pos_start, pos_end)
        self.get_psd()
        self.update_time_curves()
        self.update_psd_curves()

    def set_sample_rate(self, sr):
        self.sample_rate = sr
        srstr = eng_notation.num_to_str(self.sample_rate)
        self.gui.sampleRateLineEdit.setText(Qt.QString("%1").arg(srstr))

    def sample_rate_changed(self):
        srstr = self.gui.sampleRateLineEdit.text().toAscii()
        self.sample_rate = eng_notation.str_to_num(srstr)
        self.set_file_pos_box(self.cur_start, self.cur_stop)
        self.get_data(self.cur_start, self.cur_stop)
        self.get_psd()
        self.update_time_curves()
        self.update_psd_curves()

    def set_file_pos_box(self, start, end):
        tstart = start / self.sample_rate
        tend = end / self.sample_rate

        self.gui.filePosStartLineEdit.setText(Qt.QString("%1").arg(start))
        self.gui.filePosStopLineEdit.setText(Qt.QString("%1").arg(end))
        self.gui.filePosLengthLineEdit.setText(Qt.QString("%1").arg(end-start))

        self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
        self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
        self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tend-tstart))

    def file_position_changed(self):
        start  = self.gui.filePosStartLineEdit.text().toInt()
        end    = self.gui.filePosStopLineEdit.text().toInt()
        if((start[1] == True) and (end[1] == True)):
            self.cur_start = start[0]
            self.cur_stop = end[0]

            tstart = self.cur_start / self.sample_rate
            tend = self.cur_stop / self.sample_rate
            self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
            self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
            
            self.get_data(self.cur_start, self.cur_stop)

            self.update_time_curves()
            self.update_psd_curves()

        # If there's a non-digit character, reset box
        else:
            self.set_file_pos_box(self.cur_start, self.cur_stop)

    def file_time_changed(self):
        tstart = self.gui.fileTimeStartLineEdit.text().toDouble()
        tstop  = self.gui.fileTimeStopLineEdit.text().toDouble()
        if((tstart[1] == True) and (tstop[1] == True)):
            self.cur_start = int(tstart[0] * self.sample_rate)
            self.cur_stop = int(tstop[0] * self.sample_rate)
            self.get_data(self.cur_start, self.cur_stop)

            self.gui.filePosStartLineEdit.setText(Qt.QString("%1").arg(self.cur_start))
            self.gui.filePosStopLineEdit.setText(Qt.QString("%1").arg(self.cur_stop))

            self.update_time_curves()
            self.update_psd_curves()
        # If there's a non-digit character, reset box
        else:
            self.set_file_pos_box(self.cur_start, self.cur_stop)

    def file_length_changed(self):
        start = self.gui.filePosStartLineEdit.text().toInt()
        length = self.gui.filePosLengthLineEdit.text().toInt()
        if((start[1] == True) and (length[1] == True)):
            self.cur_start = start[0]
            self.block_length = length[0]
            self.cur_stop = self.cur_start + self.block_length

            tstart = self.cur_start / self.sample_rate
            tend = self.cur_stop / self.sample_rate
            tlen = self.block_length / self.sample_rate
            self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
            self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
            self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tlen))

            self.gui.plotHBar.setPageStep(self.block_length)

            self.get_data(self.cur_start, self.cur_stop)
            self.get_psd()

            self.update_time_curves()
            self.update_psd_curves()
        # If there's a non-digit character, reset box
        else:
            self.set_file_pos_box(self.cur_start, self.cur_stop)

    def file_time_length_changed(self):
        tstart = self.gui.fileTimeStartLineEdit.text().toDouble()
        tlength = self.gui.fileTimeLengthLineEdit.text().toDouble()
        if((tstart[1] == True) and (tlength[1] == True)):
            self.cur_start = int(tstart[0] * self.sample_rate)
            self.block_length = int(tlength[0] * self.sample_rate)
            self.cur_stop = self.cur_start + self.block_length

            tstart = self.cur_start / self.sample_rate
            tend = self.cur_stop / self.sample_rate
            tlen = self.block_length / self.sample_rate
            self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
            self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
            self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tlen))

            self.get_data(self.cur_start, self.cur_stop)
            self.get_psd()

            self.update_time_curves()
            self.update_psd_curves()
        # If there's a non-digit character, reset box
        else:
            self.set_file_pos_box(self.cur_start, self.cur_stop)


    def update_time_curves(self):
        self.icurve.setData(self.time, self.iq.imag)
        self.rcurve.setData(self.time, self.iq.real)

        # Reset the x-axis to the new time scale
        iqmax = 1.5 * max(max(self.iq.real), max(self.iq.imag))
        iqmin = 1.5 * min(min(self.iq.real), min(self.iq.imag))
        self.gui.timePlot.setAxisScale(self.gui.timePlot.xBottom,
                                       min(self.time),
                                       max(self.time))
        self.gui.timePlot.setAxisScale(self.gui.timePlot.yLeft,
                                       iqmin,
                                       iqmax)

        # Set the zoomer base to unzoom to the new axis
        self.timeZoomer.setZoomBase()
    
        self.gui.timePlot.replot()
        
    def update_psd_curves(self):
        self.psdcurve.setData(self.freq, self.iq_psd)

        self.gui.freqPlot.setAxisScale(self.gui.freqPlot.xBottom,
                                       min(self.freq),
                                       max(self.freq))

        # Set the zoomer base to unzoom to the new axis
        self.freqZoomer.setZoomBase()

        self.gui.freqPlot.replot()

    def tabChanged(self, index):
        self.gui.timePlot.replot()
        self.gui.freqPlot.replot()

    def color_black_on_white(self):
        blue = QtGui.qRgb(0x00, 0x00, 0xFF)
        red = QtGui.qRgb(0xFF, 0x00, 0x00)

        blackBrush = Qt.QBrush(Qt.QColor("black"))
        blueBrush = Qt.QBrush(Qt.QColor(blue))
        redBrush = Qt.QBrush(Qt.QColor(red))

        self.gui.timePlot.setCanvasBackground(Qt.QColor("white"))
        self.gui.freqPlot.setCanvasBackground(Qt.QColor("white"))
        self.picker.setTrackerPen(Qt.QPen(blackBrush, 2))
        self.timeZoomer.setTrackerPen(Qt.QPen(blackBrush, 2))
        self.timeZoomer.setRubberBandPen(Qt.QPen(blackBrush, 2))
        self.freqZoomer.setTrackerPen(Qt.QPen(blackBrush, 2))
        self.freqZoomer.setRubberBandPen(Qt.QPen(blackBrush, 2))
        self.psdcurve.setPen(Qt.QPen(blueBrush, 1))
        self.rcurve.setPen(Qt.QPen(blueBrush, 2))
        self.icurve.setPen(Qt.QPen(redBrush, 2))

        self.gui.timePlot.replot()
        self.gui.freqPlot.replot()

    def color_white_on_black(self):
        white = QtGui.qRgb(0xFF, 0xFF, 0xFF)
        red = QtGui.qRgb(0xFF, 0x00, 0x00)

        whiteBrush = Qt.QBrush(Qt.QColor("white"))
        whiteBrush = Qt.QBrush(Qt.QColor(white))
        redBrush = Qt.QBrush(Qt.QColor(red))
        
        self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
        self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
        self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.psdcurve.setPen(Qt.QPen(whiteBrush, 1))
        self.rcurve.setPen(Qt.QPen(whiteBrush, 2))
        self.icurve.setPen(Qt.QPen(redBrush, 2))

        self.gui.timePlot.replot()
        self.gui.freqPlot.replot()


    def color_green_on_black(self):
        green = QtGui.qRgb(0x00, 0xFF, 0x00)
        red = QtGui.qRgb(0xFF, 0x00, 0x50)

        whiteBrush = Qt.QBrush(Qt.QColor("white"))
        greenBrush = Qt.QBrush(Qt.QColor(green))
        redBrush = Qt.QBrush(Qt.QColor(red))
        
        self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
        self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
        self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.psdcurve.setPen(Qt.QPen(greenBrush, 1))
        self.rcurve.setPen(Qt.QPen(greenBrush, 2))
        self.icurve.setPen(Qt.QPen(redBrush, 2))

        self.gui.timePlot.replot()
        self.gui.freqPlot.replot()

    def color_blue_on_black(self):
        blue = QtGui.qRgb(0x00, 0x00, 0xFF)
        red = QtGui.qRgb(0xFF, 0x00, 0x00)

        whiteBrush = Qt.QBrush(Qt.QColor("white"))
        blueBrush = Qt.QBrush(Qt.QColor(blue))
        redBrush = Qt.QBrush(Qt.QColor(red))
        
        self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
        self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
        self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
        self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
        self.psdcurve.setPen(Qt.QPen(blueBrush, 1))
        self.rcurve.setPen(Qt.QPen(blueBrush, 2))
        self.icurve.setPen(Qt.QPen(redBrush, 2))

        self.gui.timePlot.replot()
        self.gui.freqPlot.replot()

def setup_options():
    usage="%prog: [options] (input_filename)"
    description = ""

    parser = OptionParser(conflict_handler="resolve", usage=usage, description=description)
    parser.add_option("-B", "--block-length", type="int", default=8192,
                      help="Specify the block size [default=%default]")
    parser.add_option("-s", "--start", type="int", default=0,
                      help="Specify where to start in the file [default=%default]")
    parser.add_option("-R", "--sample-rate", type="float", default=1.0,
                      help="Set the sampler rate of the data [default=%default]")
    parser.add_option("", "--psd-size", type="int", default=2048,
                      help="Set the size of the PSD FFT [default=%default]")
    parser.add_option("", "--spec-size", type="int", default=256,
                      help="Set the size of the spectrogram FFT [default=%default]")

    return parser

def main(args):
    parser = setup_options()
    (options, args) = parser.parse_args ()

    if(len(args) == 1):
        filename = args[0]
    else:
        filename = None

    app = Qt.QApplication(args)
    gplt = gr_plot_qt(app, filename, options)
    app.exec_()

if __name__ == '__main__':
    main(sys.argv)