# Python Workbench prototype
# Simple graphics output types
# (C)Sky Coyote, June 2007

# Import required code
import sys
from PyQt4 import QtCore, QtGui
from numpy import *
import pyfits

# Simple 1d plot window
class PlotWindow(QtGui.QWidget):
   def __init__(self, title, data):
      # Create top-level Qt window
      QtGui.QWidget.__init__(self, None)
      # Make copy of data
      self.data = data[:]
      self.setWindowTitle(self.tr(title))
      self.resize(500, 400)
     
   # Redraw
   def paintEvent(self, event):
      # Get graphics context
      painter = QtGui.QPainter()
      painter.begin(self)
      painter.setPen(QtGui.QColor(0, 0, 0))
      painter.setBrush(QtGui.QBrush())
      
      # Draw axes
      painter.drawLine(QtCore.QLineF(10.0, 10.0, 10.0, self.height() - 10.0));
      painter.drawLine(QtCore.QLineF(10.0, self.height() - 10.0, self.width() - 10.0, self.height() - 10.0));
      
      # Draw 1 curve
      if len(self.data) > 1:
         ymax = float(max(self.data))
         ymin = float(min(self.data))
         painter.setPen(QtGui.QColor(255, 0, 0))
         for i in range(len(self.data)):
            h = 10.0 + float(i) * (self.width() - 20.0) / (len(self.data) - 1.0)
            if ymax == ymin:
               v = self.height() - 10.0 - (self.height() - 20.0) / 2.0
            else:
               v = self.height() - 10.0 - (self.height() - 20.0) * (self.data[i] - ymin) / (ymax - ymin)
            if i > 0: painter.drawLine(QtCore.QLineF(h2, v2, h, v))
            h2 = h
            v2 = v

      painter.end()
      
   # Update data & redraw
   def setData(self, data):
      self.data = data[:]
      self.update()

# Display an image (QImage or numpy ndarray)
class ImageWindow(QtGui.QWidget):
   def __init__(self, title, data, min=0, max=0):
      QtGui.QWidget.__init__(self, None)
      self.data = data
      # Create pixmap from data
      self.pixmap = self.makePixmap(self.data, min, max)
      self.setWindowTitle(self.tr(title))
      self.resize(self.pixmap.width(), self.pixmap.height())
     
   # Redraw
   def paintEvent(self, event):
      painter = QtGui.QPainter()
      painter.begin(self)
      painter.setPen(QtGui.QColor(0, 0, 0))
      painter.setBrush(QtGui.QBrush())
      # Draw pixmap
      painter.drawPixmap(QtCore.QPointF(0.0, 0.0), self.pixmap)
      painter.end()
      
   # Display a different image
   def setData(self, data):
      self.data = data
      self.pixmap = self.makePixmap(self.data)
      self.update()
      
   # Create pixmap from data
   def makePixmap(self, data, min, max):
      if isinstance(data, QtGui.QImage):
         return QtGui.QPixmap.fromImage(data)
      if isinstance(data, ndarray):
         # Flip array
         data = data[::-1, :]
         # Clip array
         temp = min
         if min == 0 and max == 0: min = data.min()
         if temp == 0 and max == 0: max = data.max()
         data = data.clip(min=min, max=max)
         min = data.min()
         max = data.max()
         # Faster Python-like way to copy and scale image into RGB buffer
         byteArray = (255 * (data - min) / (max - min)).astype('b').repeat(4)
         # Pixmap data must have a Python reference, or it will be deleted!
         byteString = byteArray.tostring()
         # Create QPixmap from byte buffer
         return QtGui.QPixmap.fromImage(QtGui.QImage(byteString, data.shape[1], \
            data.shape[0], QtGui.QImage.Format_RGB32))
            
# Display an image with a circular overlay
class ImageCircleWindow(ImageWindow):
   def __init__(self, title, data, min=0, max=0, x=100, y=100, r=100):
      ImageWindow.__init__(self, title, data, min, max)
      # Set initial coords
      self.x = x
      self.y = y
      self.r = r
     
   # Redraw
   def paintEvent(self, event):
      painter = QtGui.QPainter()
      painter.begin(self)
      painter.setPen(QtGui.QColor(255, 0, 0))
      painter.setBrush(QtGui.QBrush())
      # Draw image
      painter.drawPixmap(QtCore.QPointF(0.0, 0.0), self.pixmap)
      # Draw circle over image
      painter.drawEllipse(QtCore.QRectF(self.x - self.r, self.y - self.r, 2.0 * self.r, 2.0 * self.r))
      painter.end()
                  
if __name__ == "__main__":
   app = QtGui.QApplication(sys.argv)
   window1 = PlotWindow("Plot", sin(arange(100) / 99.0 * 2.0 * pi))
   window1.show()
   window2 = ImageWindow("Image", QtGui.QImage('images/princesswands.jpg'))
   window2.show()
   window3 = ImageWindow("Image", pyfits.getdata('images/c1230.en.fits'))
   window3.show()
   window4 = ImageCircleWindow("Image", pyfits.getdata('images/c1230.en.fits'), 0, 0, 201.0, 400.0, 336.0)
   window4.show()
   sys.exit(app.exec_())