"""Read and plot a FITS image"""

# Import modules used
import sys
import pyfits
import numpy
import wx

def fits2image(fitsArray, min=-1, max=-1):
    """Convert ndarray to wx.Image"""
    # Flip image
    print 'Flipping array'
    fitsArray = fitsArray[::-1, :]
    # Clip image
    print 'Clipping array'
    if min == -1: min = fitsArray.min()
    if max == -1: max = fitsArray.max()
    fitsArray = fitsArray.clip(min=min, max=max)
    min = fitsArray.min()
    max = fitsArray.max()
    # Faster Python-like way to copy and scale image into RGB buffer
    print 'Filling RGB buffer'
    byteArray = (255 * (fitsArray - min) / (max - min)).astype('b').repeat(3)
    # Create wx.Image from byte buffer
    print 'Creating image from buffer'
    return wx.ImageFromData(fitsArray.shape[1], fitsArray.shape[0], byteArray.tostring())

class FITSWindow(wx.Window):
    """Create window which draws FITS image"""
    def __init__(self, parent, image):
        wx.Window.__init__(self, parent)
        # Convert wx.Image to bitmap
        'Converting image to bitmap'
        self.bm = image.ConvertToBitmap()
        # Bind the Paint event
        self.Bind(wx.EVT_PAINT, self.OnPaint)

    def OnPaint(self, evt):
        # create and clear the DC
        dc = wx.PaintDC(self)
        dc.Clear()
        # draw the bitmap
        print 'Drawing bitmap'
        dc.DrawBitmap(self.bm, 0, 0, True)

class FITSFrame(wx.Frame):
    """Create frame containing FITS window"""
    def __init__(self, fname, min=-1, max=-1):
        print 'Reading ', fname
        fitsArray = pyfits.getdata(fname)
        print 'Converting ndarray to image'
        img = fits2image(fitsArray, min, max)
        wx.Frame.__init__(self, parent=None, title=fname, size=(fitsArray.shape[1], fitsArray.shape[0] + 22))
        print 'Creating window'
        win = FITSWindow(self, img)

# If run as main program, read and show FITS image
if __name__ == '__main__':  
    # Create app and frame with FITS image
    print 'Creating app'
    app = wx.PySimpleApp(False)
    # Get scale from cmd line
    min = -1
    if len(sys.argv) > 2: min = float(sys.argv[2])
    max = -1
    if len(sys.argv) > 3: max = float(sys.argv[3])
    print 'Creating frame'
    frm = FITSFrame(sys.argv[1], min, max)
    print 'Showing frame'
    frm.Show()
    # Handle events
    print 'Running main loop'
    app.MainLoop()