#!/usr/bin/env python

#

# Copyright 2012 Free Software Foundation, Inc.

#

# This file is part of GNU Radio

#

# GNU Radio is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 3, or (at your option)

# any later version.

#

# GNU Radio 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 General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with GNU Radio; see the file COPYING.  If not, write to

# the Free Software Foundation, Inc., 51 Franklin Street,

# Boston, MA 02110-1301, USA.

#

from gnuradio import gr

from gnuradio.eng_option import eng_option

from optparse import OptionParser

import os, sys

try:

    from gnuradio import qtgui

    from PyQt4 import QtGui, QtCore

    import sip

except ImportError:

    print "Error: Program requires PyQt4 and gr-qtgui."

    sys.exit(1)

try:

    import scipy

except ImportError:

    print "Error: Scipy required (www.scipy.org)."

    sys.exit(1)

from plot_form import *

def read_samples_and_pad(filename, start, in_size, min_size):

    # Read in_size number of samples from file

    fhandle = open(filename, 'r')

    fhandle.seek(start*gr.sizeof_gr_complex, 0)

    data = scipy.fromfile(fhandle, dtype=scipy.complex64, count=in_size)

    data = data.tolist()

    fhandle.close()

    # If we have to, append 0's to create min_size samples of data

    if(len(data) < min_size):

        data += (min_size - len(data)) * [complex(0,0)]

    return data

class my_top_block(gr.top_block):

    def __init__(self, filelist, fc, samp_rate, psdsize, start,

                 nsamples, max_nsamples, avg=1.0):

        gr.top_block.__init__(self)

        self._filelist = filelist

        self._center_freq = fc

        self._samp_rate = samp_rate

        self._psd_size = psdsize

        self._start = start

        self._max_nsamps = max_nsamples

        self._nsigs = len(self._filelist)

        self._avg = avg

        if(nsamples is None):

            self._nsamps = max_nsamples

        else:

            self._nsamps = nsamples

        self.qapp = QtGui.QApplication(sys.argv)

        self.gui_snk = qtgui.freq_sink_c(self._psd_size, gr.firdes.WIN_BLACKMAN_hARRIS,

                                         self._center_freq, self._samp_rate,

                                         "GNU Radio PSD Plot", self._nsigs)

        n = 0

        self.srcs = list()

        for f in filelist:

            data = read_samples_and_pad(f, self._start,

                                        self._nsamps, self._psd_size)

            self.srcs.append(gr.vector_source_c(data))

            # Set default labels based on file names

            self.gui_snk.set_title(n, "{0}".format(f))

            n += 1

        self.connect(self.srcs[0], (self.gui_snk, 0))

        for i,s in enumerate(self.srcs[1:]):

            self.connect(s, (self.gui_snk, i+1))

        self.gui_snk.set_update_time(0);

        self.gui_snk.set_fft_average(self._avg)

        # Get Python Qt references

        pyQt = self.gui_snk.pyqwidget()

        self.pyWin = sip.wrapinstance(pyQt, QtGui.QWidget)

    def get_gui(self):

        return self.pyWin

    def reset(self, newstart, newnsamps):

        self.stop()

        self.wait()

        self._start = newstart

        self._nsamps = newnsamps

        for s,f in zip(self.srcs, self._filelist):

            data = read_samples_and_pad(f, self._start,

                                        self._nsamps, self._psd_size)

            s.set_data(data)

        self.start()

def main():

    description = "Plots the PSDs of a list of files."

    parser = OptionParser(option_class=eng_option, description=description,

                          conflict_handler="resolve")

    parser.add_option("-N", "--nsamples", type="int", default=None,

                      help="Set the number of samples to display [default=prints entire file]")

    parser.add_option("-S", "--start", type="int", default=0,

                      help="Starting sample number [default=%default]")

    parser.add_option("-L", "--psd-size", type="int", default=2048,

                      help="Set the FFT size of the PSD [default=%default]")

    parser.add_option("-f", "--center-frequency", type="eng_float", default=0.0,

                      help="Set the center frequency of the signal [default=%default]")

    parser.add_option("-r", "--sample-rate", type="eng_float", default=1.0,

                      help="Set the sample rate of the signal [default=%default]")

    parser.add_option("-a", "--average", type="float", default=1.0,

                      help="Set amount of averaging (smaller=more averaging) [default=%default]")

    (options, args) = parser.parse_args()

    if(len(args) < 1):

        parser.print_help()

        sys.exit(0)

    filelist = list(args)

    nsamples = options.nsamples

    # Find the smallest number of samples in all files and use that as

    # a maximum value possible.

    filesizes = []

    for f in filelist:

        if(os.path.exists(f)):

            filesizes.append(os.path.getsize(f) / gr.sizeof_gr_complex)

    max_nsamples = min(filesizes)

    tb = my_top_block(filelist,

                      options.center_frequency, options.sample_rate,

                      options.psd_size,

                      options.start, nsamples, max_nsamples,

                      options.average);

    main_box = dialog_box(tb, 'GNU Radio PSD Plot')

    main_box.show()

    tb.run()

    tb.qapp.exec_()

if __name__ == "__main__":

    try:

        main()

    except KeyboardInterrupt:

        pass