"""Start waterfall program from scratch."""
from PyQt4.QtCore import *
from PyQt4.QtGui import *
import sys
import numpy as np
import re
import threading
import Queue
import os

# ==================== some general classes ==============================

# ----------------------- class MainArea ---------------------------------
class MainArea(QFrame):
  """This is the main area of a window in which the widgets are
  laid out"""
  def __init__(self, parent):
    QFrame.__init__(self, parent)
 
# ------------------- class ResizableMainWindow --------------------------
class ResizableMainWindow(QMainWindow):
  """This is a window that can be resized; used for graphics"""
  def resizeEvent(self,event):
    self.emit(SIGNAL("resize"),event)

# ----------------------- class FigureWindow -----------------------------
class FigureWindow():
  """FigureWindow:
  This creates a resizable window with a canvas and a toolbar.
  It uses the Qt4 backend.
  Main elements of the class:
    canvas           : The FigureCanvas instance
    toolbar          : The qt.QToolBar
    window           : The qt.QMainWindow
  """
  def __init__(self,xsize=6,ysize=4):
    """This creates the figure, defines a canvas for the figure"""
    
    # create an instance of the matplotlib.figure.Figure class
    # which is a container Artist
    self.figure = matplotlib.figure.Figure(figsize=(xsize,ysize))
    # tell figure to use the Qt4 backend for the canvas
    self.canvas = FigureCanvasQTAgg(self.figure)
    
    # Now add a resizable MainWindow, which is derived from
    # PyQt4.QtGui.QMainWindow
    self.window = ResizableMainWindow()
    # This is the path to
    # /usr/share/matplotlib/mpl-data/images/matplotlib.png
    mpl_image = os.path.join( matplotlib.rcParams['datapath'],
          'images',
          'matplotlib.png' )
    # this sets the icon image in the top left
    self.window.setWindowIcon(QIcon( mpl_image ))

    # Give the keyboard focus to the figure instead of the manager
    self.canvas.setFocusPolicy( Qt.ClickFocus )
    self.canvas.setFocus()

    # connect a signal to an action, in this case, when the "destroy window"
    #button is pressed
    QObject.connect( self.window, SIGNAL( 'destroyed()' ),
                            self._widgetclosed )
    self.window._destroying = False

    # request a toolbar from the backend.  This can only be done
    # when the associated canvas and window exist.
    # _get_toolbar is defined below.
    self.toolbar = self._get_toolbar(self.canvas, self.window)
    # add the toolbar to the window
    self.window.addToolBar(self.toolbar)
    # Any messages from the toolbar are displayed in the window
    # status bar.
    QObject.connect(self.toolbar, SIGNAL("message"),
                self.window.statusBar().showMessage)

    # make the canvas the central widget of the window widget
    self.window.setCentralWidget(self.canvas)

    self.window.show()

    # attach a show method to the figure for pylab ease of use
    self.canvas.figure.show = lambda *args: self.window.show()

    def notify_axes_change( fig ):
      """Whenever the axes are changed, update the toolbar"""
      if self.toolbar != None: self.toolbar.update()

    # add an ArtistObserver to watch for changes in the axes
    self.canvas.figure.add_axobserver( notify_axes_change )

    self.window.closeEvent = self._closeH
    self.window.connect(self.window,SIGNAL("resize"),self._resizeH)

  def _closeH(self,event):
    event.ignore()

  def _resizeH(self,event):
    """A resize event causes the background to be saved."""
    self.saveBackground()

  def saveBackground(self):
    """This saves the canvas background."""
    self.canvas.draw()
    self.bbox = self.figure.bbox
    self.background = self.canvas.copy_from_bbox(self.bbox)

  def restoreBackground(self):
    self.canvas.restore_region(self.background)

  def blit(self):
    self.canvas.blit(self.bbox)

  def _widgetclosed( self ):
    """If _destroying is True, exit; otherwise, set it to True.
    It is set to False when the window is created."""
    if diag:
      print "widget closed"
    if self.window._destroying: return
    self.window._destroying = True

  def _get_toolbar(self, canvas, parent):
    """This creates a toolbar.  It must be initiated after the window,
    drawingArea and figure attributess are set."""
    toolbar = \
          matplotlib.backends.backend_qt4.NavigationToolbar2QT(canvas,
            parent, False)
    return toolbar

  def resize(self, width, height):
    'set the canvas size in pixels'
    self.window.resize(width, height)

  def destroy( self, *args ):
    if self.window._destroying: return
    self.window._destroying = True
    QObject.disconnect( self.window, SIGNAL( 'destroyed()' ),
                                   self._widgetclosed )
    if self.toolbar: self.toolbar.destroy()
    self.window.close()

  def set_window_title(self, title):
    self.window.setWindowTitle(title)

# ========================= specific classes =============================

# ----------------------- class ControlsWindow ---------------------------
class ControlsWindow(QMainWindow):
  """
  This creates the application top window with the following properties:
  menubar        - pull-down menus
  central widget - area for buttons and data values
  statusbar      - area for messages
  Class attributes are shared by all members of the class.  Note that
  ControlsWindow.view is not the same as self.view. These are configuration
  parameters with their default values if applicable:
  view             - number of spectra visible in the waterfall plot (512)
  last_read        - number of the last read spectrum, 1 .... history
  nch              - number of channels in the spectra (1024)
  oldest_displayed - spectrum at bottom of waterfall plot next to
                     disappear (0)
  history          - number of 1D spectra stored in a 2D numpy array
                     (16384)
  """
  # These are the class attributes

  def __init__(self,parent=None,view=512,nch=1024,history=16384,channel=0):
    """This creates a window with a menubar, a main area for widgets,
    and a statusbar."""
    QMainWindow.__init__(self,parent)
    self.setWindowTitle('Controls')

    # optional parameters
    self.view = view
    self.nch = nch
    self.history = history

    # initial values
    self.data_fd = None
    self.last_read = self.view # initial value
    self.oldest_displayed = 0
    # One row smaller so first row append gives right shape
    self.data = np.zeros((view-1,nch)) # spec chan 0 = average time value
    self.waterfall_on = False # initial value
    self.header = {}

    # to create a menubar, first create the actions
    
    # ....................... file menu actions .........................
    # open a data file
    datafile = QAction(QIcon(
      '/usr/share/doc/python-qt4-doc/examples/mainwindows/application/images/open.png')
      , 'Open', self)
    datafile.setShortcut('Ctrl+O')
    datafile.setStatusTip('Open new File')
    self.connect(datafile, SIGNAL('triggered()'), self.openDataFile)
    # exit program
    exitIcon = "/usr/share/python-wxglade/icons/gtk/exit.xpm"
    exit = QAction(QIcon(exitIcon), 'Exit', self)
    exit.setShortcut('Ctrl+Q')
    exit.setStatusTip('Exit application')
    self.connect(exit, SIGNAL('triggered()'), qApp, SLOT('quit()'))

    # ...................... config menu actions ............................
    # set view window range
    set_view = QAction(QIcon(), 'Set View', self)
    set_view.setStatusTip('Set time window for waterfall plot')
    self.connect(set_view, SIGNAL('triggered()'), self.setview)
    
    # create the menubar...
    # ...with a File button...
    menubar = self.menuBar()
    file = menubar.addMenu('&File')
    file.addAction(datafile)
    file.addAction(exit)
    # ... and a Config button
    config = menubar.addMenu('&Config')
    config.addAction(set_view)

    # ................. Create the main window widgets.....................
    main = MainArea(self)
    self.setCentralWidget(main)

    # Datafile
    file_lbl = QLabel(main)
    file_lbl.setText('File:')
    self.datafile =QLabel(main)
    self.datafile.setText('')
    
    # Set scan range view
    view_lbl = QLabel(main)
    view_lbl.setText('View:')
    self.view_val = QLabel(main)
    self.view_val.setText(str(view))

    # Exit program
    quit = QPushButton('Close', main)
    quit.setGeometry(10, 40, 60, 35)
    self.connect(quit, SIGNAL('clicked()'), qApp, SLOT('quit()'))
    
    # create a grid to hold the widgets and add the widgets
    grid = QGridLayout(main)
    grid.setSpacing(10)
    grid.addWidget(file_lbl,1,0)
    grid.addWidget(self.datafile,1,1)
    grid.addWidget(view_lbl, 2, 0)
    grid.addWidget(self.view_val, 2, 1)
    grid.addWidget(quit,     3, 0)
    main.setLayout(grid)
    
    # add a status for status messages
    self.statusBar()
    self.show()

  def setview(self):
    """This allows the view range to be changed from the Config menu"""
    text, ok = QInputDialog.getText(self, 'Configuration',
        'Range of spectra in view:')
    if ok:
      view = int(text)
      self.view_val.setText(str(view))
      if self.last_read >= view:
        # All spectra to be displayed are in the history buffer
        self.oldest_displayed = self.last_read - view
      else:
        # Need to prepend the zeros to the data to make up 'view' scans
        prepend = view - self.last_read
        self.data = np.append(np.zeros((prepend,self.nch)), \
                                     self.data,0)
        self.oldest_displayed = 0
        self.last_read += prepend
      self.waterfall_on = False

  def openDataFile(self):
    filename = QFileDialog.getOpenFileName(self, 'Open file','.')
    print filename
    head,tail = os.path.split(filename)
    print head,tail
    self.datafile.setText(os.path.basename(filename))
    self.data_fd = open(filename)
    self.header = self.process_header(self.data_fd)
    if self.waterfall_on:
      # OK to start a thread to read the data
      pass

  def process_header(self,fd):
    header = {}
    doing_header = True
    while(doing_header):
      line = fd.readline().strip()
      if re.search('::',line):
        [k,v] = line.split('::')
        key = k.strip()
        print key,v
        if re.search('\.',v) == None:
          header[key] = int(v.strip())
        else:
          header[key] = float(v.strip())
      elif re.search('HEADER_END',line):
        doing_header = False
    return header
    
  def closeEvent(self, event):
    """This verifies the intent to close the window and exit the program"""
    reply = QMessageBox.question(self, 'Confirm',
      "Are you sure to quit?", QMessageBox.Yes, QMessageBox.No)
    if reply == QMessageBox.Yes:
      event.accept()
    else:
      event.ignore()

class SpectrumWindow(FigureWindow,ControlsWindow):
  """This provides a waterfall plot in a FigureWindow. It needs to have
  a data array of at least 'view' by 'nch' size to fill the waterfall
  image. The window has two subplots, the upper one with a 2D line plot
  of the spectrum being added at the top of the waterfall plot and the
  bottom one with the waterfall plot. 'xlabel' and 'ylabel' refer to the
  waterfall plot."""

  def __init__(self,xsize=6,ysize=4,xlabel="",ylabel=""):
    """Creates the windowof size 'xsize' x 'ysize' and x and y labels
    defaulting to empty strings."""
    FigureWindow.__init__(self,xsize=6,ysize=4)
    self.specWin = self.window
    self.specFig = self.figure
    # Add the upper X-Y axes area to the figure for the current spectrum
    # x,y = 0.1,0.75, width = 0.8, height = 0.2.  This method also
    # creates the XY-plot
    self.make_spectrum_axes(.1,.75,.75,.2)
    # Add a lower axes for the waterfall image. The bottom axis goes from
    # x,y = 0.1,0.75, width 0.8, height = 0.7.  This method also creates
    # the image.
    self.make_waterfall_axes(xlabel,ylabel,.1,.05,.75,.65)
    # add a colorbar on the right
    self.colorbar_axes = self.specFig.add_axes([0.85,0.05,.15,.9],aspect=20)
    self.specFig.colorbar(self.waterfall_artist_bottom,
                          cax=self.colorbar_axes)
    
    self.specFig.canvas.draw()
    self.window.update()
    self.window.repaint()
    # This defines the bbox used by 'blit'
    self.saveBackground()

  
  def make_spectrum_axes(self,llx,lly,width,height):
    """This creates the axes for the spectrum (top axes)"""
    self.waterfall_top_axes    = self.specFig.add_axes([llx,lly,width,height])
    # Plot the current spectrum starting
    if diag:
      print "open_waterfall: View =",view
      print "open_waterfall: Data shape =",ControlsWindow.data.shape
    ydata = ControlsWindow.data[view-1,:]
    self.max_data = np.amax(ydata)
    self.water_spectrum_lines = \
        self.waterfall_top_axes.plot(ydata,animated=True)
    # We'd like the ticks across the top but this doesn't seen to do it
    self.waterfall_top_axes.set_axisbelow(False)
    self.waterfall_top_axes.set_xticks([])
    # set the limits to channel numbers
    self.waterfall_top_axes.set_xlim(0,ControlsWindow.nch-1)
    self.waterfall_top_axes.set_ylim(0,self.max_data)
    # When hold is True, subsequent plot commands will add to the
    # current axes.  If it is None, the hold state is toggled. In
    # this case each plot commands replaces the existing line.
    self.waterfall_top_axes.hold(False)
  
  def make_waterfall_axes(self,xlabel,ylabel,llx,lly,width,height):
    """This creates the axes for the waterfall plot (bottom axes).
    It shares the X-axis with the top one.  This means that scrolling
    and zooming along the x-axis affects both plots."""
    self.waterfall_bottom_axes = self.specFig.add_axes([llx,lly,width,height],
        sharex = self.waterfall_top_axes)
    # show the initial view into the data array
    self.display_waterfall()
    # add the labels
    self.waterfall_bottom_axes.set_xlabel(xlabel)
    self.waterfall_bottom_axes.set_ylabel(ylabel)
    # The waterfall plot is now on.
    ControlsWindow.waterfall_on = True

  def display_waterfall(self):
    """This creates the waterfall plot"""
    # clear the bottom axes
    self.waterfall_bottom_axes.cla()
    im = ControlsWindow.data[
      ControlsWindow.oldest_displayed:
      ControlsWindow.oldest_displayed + view,
      :]
    self.waterfall_artist_bottom = \
        self.waterfall_bottom_axes.imshow(
          im,
          aspect = 'auto',
          cmap=matplotlib.cm.spectral,
          interpolation='bicubic',
          animated = True,
          origin='lower',
          extent=(0,
                  ControlsWindow.nch-1,
                  ControlsWindow.last_read,
                  ControlsWindow.last_read-view))
    # set the limits to channel number
    self.waterfall_bottom_axes.set_xlim(0,ControlsWindow.nch-1)
    # set the view range
    self.waterfall_bottom_axes.set_ylim(
        ControlsWindow.last_read,
        ControlsWindow.last_read-view)
    self.waterfall_bottom_axes.autoscale_view()

  def update(self):
    """This is called first whenever a new spectrum has been appended
    to the data array.   The data array has up to 'history' number of rows.
    If it is smaller than that, this keeps adding the spectra to the
    bottom of the array.  Once 'history' rows exist, a spectrum is
    removed from the top for every one added to the bottom."""
    if diag:
      print "update_waterfall: Data shape =",ControlsWindow.data.shape
    #if ControlsWindow.last_read == self.history:
      # The history buffer is full so start removing spectra from the
      # top of the buffer. An new spectrum has already... to come
    # get the new waterfall image
    im = ControlsWindow.data[ControlsWindow.oldest_displayed: \
                     ControlsWindow.oldest_displayed+view,:]
    # Just update the image data
    self.waterfall_artist_bottom.set_data(im)
    self.waterfall_bottom_axes.draw_artist(self.waterfall_artist_bottom)
    # update the spectrum
    if diag:
      print "update_waterfall: spectrum index, last_read =",\
          ControlsWindow.oldest_displayed+view,\
          ControlsWindow.last_read
    ydata = ControlsWindow.data[
            ControlsWindow.oldest_displayed+view-1,:]
    self.water_spectrum_lines[0].set_ydata(ydata)
    self.waterfall_top_axes.redraw_in_frame()
    self.waterfall_top_axes.draw_artist(self.water_spectrum_lines[0])

# ========================= thread functions ===========================

def process_data_line(line,nchan):
  data = line.split()
  mean = []
  rms = []
  skew = []
  kurtosis = []
  sec = int(data[0])
  ms = int(data[1])
  for i in range(2,2+4*nchan,4):
    mean.append(float(data[i]))
    rms.append(float(data[i+1]))
    kurtosis.append(float(data[i+2]))
    skew.append(float(data[i+3]))
  return (sec,ms,mean,rms,kurtosis,skew)

def get_new_data():
  counter = 0
  while(counter<1000):
    line = fd.readline().strip()
    if line != '':
      sec,ms,mean,rms,kurt,skew = process_data_line(line,head['NOCHAN'])
      if counter == 0:
        # initialize the arrays
        means = array(mean,ndmin=2)
        rootmeansquare = array(rms,ndmin=2)
        kurtosis = array(kurt,ndmin=2)
        skewness = array(skew,ndmin=2)
      else:
        # append to the arrays
        means = append(means,array(mean,ndmin=2),axis=0)
        rootmeansquare = append(rootmeansquare,array(rms,ndmin=2),axis=0)
        kurtosis = append(kurtosis,array(kurt,ndmin=2),axis=0)
        skewness = append(skewness,array(skew,ndmin=2),axis=0)
      counter += 1
  # Pad the arrays
  if diag:
    print "Kurtosis array shape:",kurtosis.shape
  means = append(means,zeros((24,head['NOCHAN'])),axis=0)
  rootmeansquare = append(rootmeansquare,zeros((24,head['NOCHAN'])),axis=0)
  kurtosis = append(kurtosis,zeros((24,head['NOCHAN'])),axis=0)
  skewness = append(skewness,zeros((24,head['NOCHAN'])),axis=0)
  # compute the kurtosis spectrum
  spectrum = []
  for i in range(head['NOCHAN']):
    if diag:
      print "Processing scan",i
    spec = rfft(kurtosis[:,i])
    spectrum.append(abs(spec))
  return spectrum


diag = True
# This creates the application object
app = QApplication(sys.argv)
controls = ControlsWindow()
data_fd = controls.data_fd
if diag:
  print "controls.nch:",controls.nch
  print "Controls Window data shape:",controls.data.shape
  print "Data file:",data_fd

# The file reading thread puts data in here:
Qin = Queue.Queue()

reader = threading.Thread(target=get_new_data)
sys.exit(app.exec_())