import os

from wxPython.wx import *

from Numeric import *

import pygame
from pygame import display,event,surfarray,time


class wxSDLWindow(wxFrame):
    def __init__(self, parent, id, title = 'SDL window', **options):
        options['style'] = wxDEFAULT_FRAME_STYLE | wxTRANSPARENT_WINDOW
        wxFrame.__init__(*(self, parent, id, title), **options)

        self._initialized = 0
        self._resized = 0
        self._surface = None
        self.__needsDrawing = 1

        EVT_IDLE(self, self.OnIdle)

    def OnIdle(self, ev):
        if not self._initialized or self._resized:
            if not self._initialized:
                # get the handle
                hwnd = self.GetHandle()
                os.putenv('SDL_WINDOWID', str(hwnd))
                pygame.init()
                EVT_SIZE(self, self.OnSize)
                self._initialized = 1
        else:
            self._resized = 0

        x,y = self.GetSizeTuple()
        self._surface = display.set_mode((x,y))

        if self.__needsDrawing:
            self.draw()

    def OnPaint(self, ev):
        self.__needsDrawing = 1

    def OnSize(self, ev):
        self._resized = 1
        ev.Skip()

    def draw(self):
        raise NotImplementedError('please define a .draw() method!')

    def getSurface(self):
        return self._surface


if __name__ == "__main__":

    class GradientWindow(wxSDLWindow):
        "simple random color gradient done with a NumPy array."
        def draw(self):
            surface = self.getSurface()
            if not surface is None:
                topcolor = 5
                bottomcolor = 100
                # topcolor,bottomcolor = random.uniform_int_sample(0, 255, (2, 3))
                diff = bottomcolor - topcolor
                width,height = surface.get_size()
                # create array from 0.0 to 1.0 triplets
                column = arange(height, typecode=Float)/height
                column = repeat(column[:, NewAxis], [3], 1)
                # create a single column of gradient
                column = topcolor + (diff * column).astype(Int)
                # make the column a 3d image column by adding X
                column = column.astype(UnsignedInt8)[NewAxis,:,:]
                #3d array into 2d array
                column = surfarray.map_array(surface, column)
                # stretch the column into a full image
                surfarray.blit_array(self._surface,
                resize(column, (width, height)))
                display.flip()

    def gradienttest():
        app = wxPySimpleApp()
        sizeT = (640,480)
        w = GradientWindow(None, -1, size = sizeT)
        w.Show(1)
        app.MainLoop()

    gradienttest()