#!/usr/bin/env python
#
# Copyright 2003,2004,2006,2007,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, gru, eng_notation
from gnuradio import analog
from gnuradio import blocks
from gnuradio import wxgui
from gnuradio.wxgui import stdgui2
import wx
import gnuradio.wxgui.plot as plot
import numpy
import struct

default_scopesink_size = (640, 240)
default_v_scale = 1000
default_frame_decim = gr.prefs().get_long('wxgui', 'frame_decim', 1)

class scope_sink_f(gr.hier_block2):
    def __init__(self, parent, title='', sample_rate=1,
                 size=default_scopesink_size, frame_decim=default_frame_decim,
                 v_scale=default_v_scale, t_scale=None, num_inputs=1, **kwargs):

        gr.hier_block2.__init__(self, "scope_sink_f",
                                gr.io_signature(num_inputs, num_inputs, gr.sizeof_float),
                                gr.io_signature(0,0,0))

        msgq = gr.msg_queue(2)         # message queue that holds at most 2 messages
        self.guts = wxgui.oscope_sink_f(sample_rate, msgq)
        for i in range(num_inputs):
          self.connect((self, i), (self.guts, i))

        self.win = scope_window(win_info(msgq, sample_rate, frame_decim,
                                          v_scale, t_scale, self.guts, title), parent)

    def set_sample_rate(self, sample_rate):
        self.guts.set_sample_rate(sample_rate)
        self.win.info.set_sample_rate(sample_rate)

class scope_sink_c(gr.hier_block2):
    def __init__(self, parent, title='', sample_rate=1,
                 size=default_scopesink_size, frame_decim=default_frame_decim,
                 v_scale=default_v_scale, t_scale=None, num_inputs=1, xy_mode=False, **kwargs):

        gr.hier_block2.__init__(self, "scope_sink_c",
                                gr.io_signature(num_inputs, num_inputs, gr.sizeof_gr_complex),
                                gr.io_signature(0,0,0))

        msgq = gr.msg_queue(2)         # message queue that holds at most 2 messages
        self.guts = wxgui.oscope_sink_f(sample_rate, msgq)
        for i in range(num_inputs):
	        c2f = blocks.complex_to_float()
	        self.connect((self, i), c2f)
	        self.connect((c2f, 0),(self.guts, 2*i+0))
	        self.connect((c2f, 1),(self.guts, 2*i+1))

        self.win = scope_window(win_info(msgq, sample_rate, frame_decim,
                                         v_scale, t_scale, self.guts, title), parent)
        self.win.info.xy = xy_mode

    def set_sample_rate(self, sample_rate):
        self.guts.set_sample_rate(sample_rate)
        self.win.info.set_sample_rate(sample_rate)

class constellation_sink(scope_sink_c):
    def __init__(self, parent, title='Constellation', sample_rate=1,
                 size=default_scopesink_size, frame_decim=default_frame_decim):
        scope_sink_c.__init__(self, parent=parent, title=title, sample_rate=sample_rate,
                              size=size, frame_decim=frame_decim)
        self.win.info.xy = True	#constellation mode

# ========================================================================


time_base_list = [                      # time / division
    1.0e-7,   # 100ns / div
    2.5e-7,
    5.0e-7,
    1.0e-6,   #   1us / div
    2.5e-6,
    5.0e-6,
    1.0e-5,   #  10us / div
    2.5e-5,
    5.0e-5,
    1.0e-4,   # 100us / div
    2.5e-4,
    5.0e-4,
    1.0e-3,   #   1ms / div
    2.5e-3,
    5.0e-3,
    1.0e-2,   #  10ms / div
    2.5e-2,
    5.0e-2
    ]

v_scale_list = [ # counts / div, LARGER gains are SMALLER /div, appear EARLIER
    2.0e-3,   # 2m / div, don't call it V/div it's actually counts/div
    5.0e-3,
    1.0e-2,
    2.0e-2,
    5.0e-2,
    1.0e-1,
    2.0e-1,
    5.0e-1,
    1.0e+0,
    2.0e+0,
    5.0e+0,
    1.0e+1,
    2.0e+1,
    5.0e+1,
    1.0e+2,
    2.0e+2,
    5.0e+2,
    1.0e+3,
    2.0e+3,
    5.0e+3,
    1.0e+4 # 10000 /div, USRP full scale is -/+ 32767
    ]


wxDATA_EVENT = wx.NewEventType()

def EVT_DATA_EVENT(win, func):
    win.Connect(-1, -1, wxDATA_EVENT, func)

class DataEvent(wx.PyEvent):
    def __init__(self, data):
        wx.PyEvent.__init__(self)
        self.SetEventType(wxDATA_EVENT)
        self.data = data

    def Clone(self):
        self.__class__(self.GetId())


class win_info(object):
    __slots__ = ['msgq', 'sample_rate', 'frame_decim', 'v_scale',
                 'scopesink', 'title',
                 'time_scale_cursor', 'v_scale_cursor', 'marker', 'xy',
                 'autorange', 'running']

    def __init__(self, msgq, sample_rate, frame_decim, v_scale, t_scale,
                 scopesink, title = "Oscilloscope", xy=False):
        self.msgq = msgq
        self.sample_rate = sample_rate
        self.frame_decim = frame_decim
        self.scopesink = scopesink
        self.title = title;

        self.time_scale_cursor = gru.seq_with_cursor(time_base_list, initial_value = t_scale)
        self.v_scale_cursor = gru.seq_with_cursor(v_scale_list, initial_value = v_scale)

        self.marker = 'line'
        self.xy = xy
        self.autorange = not v_scale
        self.running = True

    def get_time_per_div(self):
        return self.time_scale_cursor.current()

    def get_volts_per_div(self):
        return self.v_scale_cursor.current()

    def set_sample_rate(self, sample_rate):
        self.sample_rate = sample_rate

    def get_sample_rate(self):
        return self.sample_rate

    def get_decimation_rate(self):
        return 1.0

    def set_marker(self, s):
        self.marker = s

    def get_marker(self):
        return self.marker


class input_watcher(gru.msgq_runner):
    def __init__(self, msgq, event_receiver, frame_decim, **kwds):
        self.event_receiver = event_receiver
        self.frame_decim = frame_decim
        self.iscan = 0
        gru.msgq_runner.__init__(self, msgq, self.handle_msg)

    def handle_msg(self, msg):
        if self.iscan == 0:            # only display at frame_decim
            self.iscan = self.frame_decim

            nchan = int(msg.arg1())    # number of channels of data in msg
            nsamples = int(msg.arg2()) # number of samples in each channel

            s = msg.to_string()      # get the body of the msg as a string

            bytes_per_chan = nsamples * gr.sizeof_float

            records = []
            for ch in range(nchan):

                start = ch * bytes_per_chan
                chan_data = s[start:start+bytes_per_chan]
                rec = numpy.fromstring(chan_data, numpy.float32)
                records.append(rec)

            # print "nrecords = %d, reclen = %d" % (len (records),nsamples)

            de = DataEvent(records)
            wx.PostEvent(self.event_receiver, de)
            records = []
            del de

        self.iscan -= 1


class scope_window(wx.Panel):

    def __init__(self, info, parent, id = -1,
                 pos = wx.DefaultPosition, size = wx.DefaultSize, name = ""):
        wx.Panel.__init__(self, parent, -1)
        self.info = info

        vbox = wx.BoxSizer(wx.VERTICAL)

        self.graph = graph_window(info, self, -1)

        vbox.Add(self.graph, 1, wx.EXPAND)
        vbox.Add(self.make_control_box(), 0, wx.EXPAND)
        vbox.Add(self.make_control2_box(), 0, wx.EXPAND)

        self.sizer = vbox
        self.SetSizer(self.sizer)
        self.SetAutoLayout(True)
        self.sizer.Fit(self)
        self.set_autorange(self.info.autorange)


    # second row of control buttons etc. appears BELOW control_box
    def make_control2_box(self):
        ctrlbox = wx.BoxSizer(wx.HORIZONTAL)

        self.inc_v_button = wx.Button(self, 1101, " < ", style=wx.BU_EXACTFIT)
        self.inc_v_button.SetToolTipString("Increase vertical range")
        wx.EVT_BUTTON(self, 1101, self.incr_v_scale) # ID matches button ID above

        self.dec_v_button  = wx.Button(self, 1100, " > ", style=wx.BU_EXACTFIT)
        self.dec_v_button.SetToolTipString("Decrease vertical range")
        wx.EVT_BUTTON(self, 1100, self.decr_v_scale)

        self.v_scale_label = wx.StaticText(self, 1002, "None") # vertical /div
        self.update_v_scale_label()

        self.autorange_checkbox = wx.CheckBox(self, 1102, "Autorange")
        self.autorange_checkbox.SetToolTipString("Select autorange on/off")
        wx.EVT_CHECKBOX(self, 1102, self.autorange_checkbox_event)

        ctrlbox.Add((5,0) ,0) # left margin space
        ctrlbox.Add(self.inc_v_button, 0, wx.EXPAND)
        ctrlbox.Add(self.dec_v_button, 0, wx.EXPAND)
        ctrlbox.Add(self.v_scale_label, 0, wx.ALIGN_CENTER)
        ctrlbox.Add((20,0) ,0) # spacer
        ctrlbox.Add(self.autorange_checkbox, 0, wx.ALIGN_CENTER)

        return ctrlbox

    def make_control_box(self):
        ctrlbox = wx.BoxSizer(wx.HORIZONTAL)

        tb_left = wx.Button(self, 1001, " < ", style=wx.BU_EXACTFIT)
        tb_left.SetToolTipString("Increase time base")
        wx.EVT_BUTTON(self, 1001, self.incr_timebase)


        tb_right = wx.Button(self, 1000, " > ", style=wx.BU_EXACTFIT)
        tb_right.SetToolTipString("Decrease time base")
        wx.EVT_BUTTON(self, 1000, self.decr_timebase)

        self.time_base_label = wx.StaticText(self, 1002, "")
        self.update_timebase_label()

        ctrlbox.Add((5,0) ,0)
        # ctrlbox.Add(wx.StaticText(self, -1, "Horiz Scale: "), 0, wx.ALIGN_CENTER)
        ctrlbox.Add(tb_left, 0, wx.EXPAND)
        ctrlbox.Add(tb_right, 0, wx.EXPAND)
        ctrlbox.Add(self.time_base_label, 0, wx.ALIGN_CENTER)

        ctrlbox.Add((10,0) ,1)            # stretchy space

        ctrlbox.Add(wx.StaticText(self, -1, "Trig: "), 0, wx.ALIGN_CENTER)
        self.trig_chan_choice = wx.Choice(self, 1004,
                                          choices = ['Ch1', 'Ch2', 'Ch3', 'Ch4'])
        self.trig_chan_choice.SetToolTipString("Select channel for trigger")
        wx.EVT_CHOICE(self, 1004, self.trig_chan_choice_event)
        ctrlbox.Add(self.trig_chan_choice, 0, wx.ALIGN_CENTER)

        self.trig_mode_choice = wx.Choice(self, 1005,
                                          choices = ['Free', 'Auto', 'Norm'])
        self.trig_mode_choice.SetSelection(1)
        self.trig_mode_choice.SetToolTipString("Select trigger slope or Auto (untriggered roll)")
        wx.EVT_CHOICE(self, 1005, self.trig_mode_choice_event)
        ctrlbox.Add(self.trig_mode_choice, 0, wx.ALIGN_CENTER)

        trig_level50 = wx.Button(self, 1006, "50%")
        trig_level50.SetToolTipString("Set trigger level to 50%")
        wx.EVT_BUTTON(self, 1006, self.set_trig_level50)
        ctrlbox.Add(trig_level50, 0, wx.EXPAND)

        run_stop = wx.Button(self, 1007, "Run/Stop")
        run_stop.SetToolTipString("Toggle Run/Stop mode")
        wx.EVT_BUTTON(self, 1007, self.run_stop)
        ctrlbox.Add(run_stop, 0, wx.EXPAND)

        ctrlbox.Add((10, 0) ,1)            # stretchy space

        ctrlbox.Add(wx.StaticText(self, -1, "Fmt: "), 0, wx.ALIGN_CENTER)
        self.marker_choice = wx.Choice(self, 1002, choices = self._marker_choices)
        self.marker_choice.SetToolTipString("Select plotting with lines, pluses or dots")
        wx.EVT_CHOICE(self, 1002, self.marker_choice_event)
        ctrlbox.Add(self.marker_choice, 0, wx.ALIGN_CENTER)

        self.xy_choice = wx.Choice(self, 1003, choices = ['X:t', 'X:Y'])
        self.xy_choice.SetToolTipString("Select X vs time or X vs Y display")
        wx.EVT_CHOICE(self, 1003, self.xy_choice_event)
        ctrlbox.Add(self.xy_choice, 0, wx.ALIGN_CENTER)

        return ctrlbox

    _marker_choices = ['line', 'plus', 'dot']

    def update_timebase_label(self):
        time_per_div = self.info.get_time_per_div()
        s = ' ' + eng_notation.num_to_str(time_per_div) + 's/div'
        self.time_base_label.SetLabel(s)

    def decr_timebase(self, evt):
        self.info.time_scale_cursor.prev()
        self.update_timebase_label()

    def incr_timebase(self, evt):
        self.info.time_scale_cursor.next()
        self.update_timebase_label()

    def update_v_scale_label(self):
        volts_per_div = self.info.get_volts_per_div()
        s = ' ' + eng_notation.num_to_str(volts_per_div) + '/div' # Not V/div
        self.v_scale_label.SetLabel(s)

    def decr_v_scale(self, evt):
        self.info.v_scale_cursor.prev()
        self.update_v_scale_label()

    def incr_v_scale(self, evt):
        self.info.v_scale_cursor.next()
        self.update_v_scale_label()

    def marker_choice_event(self, evt):
        s = evt.GetString()
        self.set_marker(s)

    def set_autorange(self, on):
        if on:
            self.v_scale_label.SetLabel(" (auto)")
            self.info.autorange = True
            self.autorange_checkbox.SetValue(True)
            self.inc_v_button.Enable(False)
            self.dec_v_button.Enable(False)
        else:
            if self.graph.y_range:
                (l,u) = self.graph.y_range # found by autorange
                self.info.v_scale_cursor.set_index_by_value((u-l)/8.0)
            self.update_v_scale_label()
            self.info.autorange = False
            self.autorange_checkbox.SetValue(False)
            self.inc_v_button.Enable(True)
            self.dec_v_button.Enable(True)

    def autorange_checkbox_event(self, evt):
        if evt.Checked():
            self.set_autorange(True)
        else:
            self.set_autorange(False)

    def set_marker(self, s):
        self.info.set_marker(s)        # set info for drawing routines
        i = self.marker_choice.FindString(s)
        assert i >= 0, "Hmmm, set_marker problem"
        self.marker_choice.SetSelection(i)

    def set_format_line(self):
        self.set_marker('line')

    def set_format_dot(self):
        self.set_marker('dot')

    def set_format_plus(self):
        self.set_marker('plus')

    def xy_choice_event(self, evt):
        s = evt.GetString()
        self.info.xy = s == 'X:Y'

    def trig_chan_choice_event(self, evt):
        s = evt.GetString()
        ch = int(s[-1]) - 1
        self.info.scopesink.set_trigger_channel(ch)

    def trig_mode_choice_event(self, evt):
        sink = self.info.scopesink
        s = evt.GetString()
        if s == 'Norm':
            sink.set_trigger_mode(wxgui.TRIG_MODE_NORM)
        elif s == 'Auto':
            sink.set_trigger_mode(wxgui.TRIG_MODE_AUTO)
        elif s == 'Free':
            sink.set_trigger_mode(wxgui.TRIG_MODE_FREE)
        else:
            assert 0, "Bad trig_mode_choice string"

    def set_trig_level50(self, evt):
        self.info.scopesink.set_trigger_level_auto()

    def run_stop(self, evt):
        self.info.running = not self.info.running


class graph_window(plot.PlotCanvas):

    channel_colors = ['BLUE', 'RED',
                      'CYAN', 'MAGENTA', 'GREEN', 'YELLOW']

    def __init__(self, info, parent, id = -1,
                  pos = wx.DefaultPosition, size = (640, 240),
                  style = wx.DEFAULT_FRAME_STYLE, name = ""):
        plot.PlotCanvas.__init__(self, parent, id, pos, size, style, name)

        self.SetXUseScopeTicks(True)
        self.SetEnableGrid(True)
        self.SetEnableZoom(True)
        self.SetEnableLegend(True)
        # self.SetBackgroundColour('black')

        self.info = info;
        self.y_range = None
        self.x_range = None
        self.avg_y_min = None
        self.avg_y_max = None
        self.avg_x_min = None
        self.avg_x_max = None

        EVT_DATA_EVENT(self, self.format_data)

        self.input_watcher = input_watcher(info.msgq, self, info.frame_decim)

    def channel_color(self, ch):
        return self.channel_colors[ch % len(self.channel_colors)]

    def format_data(self, evt):
        if not self.info.running:
            return

        if self.info.xy:
            self.format_xy_data(evt)
            return

        info = self.info
        records = evt.data
        nchannels = len(records)
        npoints = len(records[0])

        objects = []

        Ts = 1.0 / (info.get_sample_rate() / info.get_decimation_rate())
        x_vals = Ts * numpy.arange(-npoints/2, npoints/2)

        # preliminary clipping based on time axis here, instead of in graphics code
        time_per_window = self.info.get_time_per_div() * 10
        n = int(time_per_window / Ts + 0.5)
        n = n & ~0x1                    # make even
        n = max(2, min(n, npoints))

        self.SetXUseScopeTicks(True)   # use 10 divisions, no labels

        for ch in range(nchannels):
            r = records[ch]

            # plot middle n points of record

            lb = npoints/2 - n/2
            ub = npoints/2 + n/2
            # points = zip(x_vals[lb:ub], r[lb:ub])
            points = numpy.zeros((ub-lb, 2), numpy.float64)
            points[:,0] = x_vals[lb:ub]
            points[:,1] = r[lb:ub]

            m = info.get_marker()
            if m == 'line':
                objects.append(plot.PolyLine(points,
                                             colour=self.channel_color(ch),
                                             legend=('Ch%d' % (ch+1,))))
            else:
                objects.append(plot.PolyMarker(points,
                                               marker=m,
                                               colour=self.channel_color(ch),
                                               legend=('Ch%d' % (ch+1,))))

        graphics = plot.PlotGraphics(objects,
                                     title=self.info.title,
                                     xLabel = '', yLabel = '')

        time_per_div = info.get_time_per_div()
        x_range = (-5.0 * time_per_div, 5.0 * time_per_div) # ranges are tuples!
        volts_per_div = info.get_volts_per_div()
        if not self.info.autorange:
            self.y_range = (-4.0 * volts_per_div, 4.0 * volts_per_div)
        self.Draw(graphics, xAxis=x_range, yAxis=self.y_range)
        self.update_y_range() # autorange to self.y_range


    def format_xy_data(self, evt):
        info = self.info
        records = evt.data
        nchannels = len(records)
        npoints = len(records[0])

        if nchannels < 2:
            return

        objects = []
        # points = zip(records[0], records[1])
        points = numpy.zeros((len(records[0]), 2), numpy.float32)
        points[:,0] = records[0]
        points[:,1] = records[1]

        self.SetXUseScopeTicks(False)

        m = info.get_marker()
        if m == 'line':
            objects.append(plot.PolyLine(points,
                                         colour=self.channel_color(0)))
        else:
            objects.append(plot.PolyMarker(points,
                                           marker=m,
                                           colour=self.channel_color(0)))

        graphics = plot.PlotGraphics(objects,
                                     title=self.info.title,
                                     xLabel = 'I', yLabel = 'Q')

        self.Draw(graphics, xAxis=self.x_range, yAxis=self.y_range)
        self.update_y_range()
        self.update_x_range()


    def update_y_range(self):
        alpha = 1.0/25
        graphics = self.last_draw[0]
        p1, p2 = graphics.boundingBox()     # min, max points of graphics

        if self.avg_y_min: # prevent vertical scale from jumping abruptly --?
            self.avg_y_min = p1[1] * alpha + self.avg_y_min * (1 - alpha)
            self.avg_y_max = p2[1] * alpha + self.avg_y_max * (1 - alpha)
        else: # initial guess
            self.avg_y_min = p1[1] # -500.0 workaround, sometimes p1 is ~ 10^35
            self.avg_y_max = p2[1] # 500.0

        self.y_range = self._axisInterval('auto', self.avg_y_min, self.avg_y_max)
        # print "p1 %s  p2 %s  y_min %s  y_max %s  y_range %s" \
        #        % (p1, p2, self.avg_y_min, self.avg_y_max, self.y_range)


    def update_x_range(self):
        alpha = 1.0/25
        graphics = self.last_draw[0]
        p1, p2 = graphics.boundingBox()     # min, max points of graphics

        if self.avg_x_min:
            self.avg_x_min = p1[0] * alpha + self.avg_x_min * (1 - alpha)
            self.avg_x_max = p2[0] * alpha + self.avg_x_max * (1 - alpha)
        else:
            self.avg_x_min = p1[0]
            self.avg_x_max = p2[0]

        self.x_range = self._axisInterval('auto', self.avg_x_min, self.avg_x_max)


# ----------------------------------------------------------------
# Stand-alone test application
# ----------------------------------------------------------------

class test_top_block(stdgui2.std_top_block):
    def __init__(self, frame, panel, vbox, argv):
        stdgui2.std_top_block.__init__(self, frame, panel, vbox, argv)

        if len(argv) > 1:
            frame_decim = int(argv[1])
        else:
            frame_decim = 1

        if len(argv) > 2:
            v_scale = float(argv[2])  # start up at this v_scale value
        else:
            v_scale = None  # start up in autorange mode, default

        if len(argv) > 3:
            t_scale = float(argv[3])  # start up at this t_scale value
        else:
            t_scale = None  # old behavior

        print "frame decim %s  v_scale %s  t_scale %s" % (frame_decim,v_scale,t_scale)

        input_rate = 1e6

        # Generate a complex sinusoid
        self.src0 = analog.sig_source_c(input_rate, analog.GR_SIN_WAVE, 25.1e3, 1e3)

        # We add this throttle block so that this demo doesn't suck down
        # all the CPU available.  You normally wouldn't use it...
        self.thr = blocks.throttle(gr.sizeof_gr_complex, input_rate)

        scope = scope_sink_c(panel,"Secret Data",sample_rate=input_rate,
                             frame_decim=frame_decim,
                             v_scale=v_scale, t_scale=t_scale)
        vbox.Add(scope.win, 1, wx.EXPAND)

        # Ultimately this will be
        # self.connect("src0 throttle scope")
	self.connect(self.src0, self.thr, scope)

def main():
    app = stdgui2.stdapp(test_top_block, "O'Scope Test App")
    app.MainLoop()

if __name__ == '__main__':
    main()

# ----------------------------------------------------------------
