"""Demonstrate a complicated dynamically-generated dataset.
"""

# __init__.py

import sys, copy, logging, sys

import Image #PIL

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
from pyglew import *

from glviewer import load_texture

import numpy

import theano
from theano.compile import shared
from theano.compile import pfunc as function

_logger = logging.getLogger('gldataset')
def debug(*msg): _logger.debug(' '.join(str(m) for m in msg))
def info(*msg): _logger.info(' '.join(str(m) for m in msg))
def warn(*msg): _logger.warn(' '.join(str(m) for m in msg))
def warning(*msg): _logger.warning(' '.join(str(m) for m in msg))
def error(*msg): _logger.error(' '.join(str(m) for m in msg))

def init_GL(shape=(64,64), title='Offscreen rendering using FB0'):
    if not init_GL.done:
        w, h = shape
        init_GL.done = True
        info('initializing OpenGl subsystem')
        glutInit (sys.argv)
        glutInitDisplayMode (GLUT_DOUBLE | GLUT_RGBA | GLUT_DEPTH)
        glutInitWindowSize (w,h)
        init_GL.window = glutCreateWindow (title)
        glewInit()

        glEnable(GL_TEXTURE_2D)
        glClearColor(0.0, 0.0, 0.0, 0.0)    # This Will Clear The Background Color To Black
        glClearDepth(1.0)                                   # Enables Clearing Of The Depth Buffer
        glDepthFunc(GL_LESS)                                # The Type Of Depth Test To Do
        glEnable(GL_DEPTH_TEST)                             # Enables Depth Testing
        glShadeModel(GL_SMOOTH)                             # Enables Smooth Color Shading
            
        #glMatrixMode(GL_PROJECTION)
        #glLoadIdentity()                                    # Reset The Projection Matrix
                                                            # Calculate The Aspect Ratio Of The Window
        #gluPerspective(45.0, float(64)/float(64), 0.1, 100.0)
        glMatrixMode(GL_MODELVIEW)
init_GL.done = False

class PBufRenderer(object):
    """Render an OpenGL program to a framebuffer instead of the screen.

    The way to use this class is to enclose all the OpenGL commands you want to render between
    a call to setup() and a call to render().  So you would render a frame like this:

    .. code-block:: python
        
        p = PBufRenderer(shape)
        p.setup()
        my_display_code()
        a = p.render()
        my_display_code()
        b = p.render()

    After running this code, 'a' and 'b' will be numpy arrays of shape `shape` + (3,) containing an
    RBG rendering of your display_code.
    """
    def __init__(self, size=(128,128), upside_down=False):
        """ Offscreen rendering
        
        Save an offscreen rendering of size (w,h) to filename.
        """
        
        def round2 (n):
            """ Get nearest power of two superior to n """
            f = 1
            while f<n:
                f*= 2
            return f

        if size == None:
            size = (512,512)
        w = round2 (size[0])
        h = round2 (size[1])

        image = Image.new ("RGB", (w, h), (0, 0, 0))
        bits = image.tostring("raw", "RGBX", 0, -1)

        debug('allocating framebuffer')
        framebuffer = glGenFramebuffersEXT (1)
        glBindFramebufferEXT(GL_FRAMEBUFFER_EXT, framebuffer)

        debug('allocating depthbuffer')
        depthbuffer = glGenRenderbuffersEXT (1)
        glBindRenderbufferEXT (GL_RENDERBUFFER_EXT,depthbuffer)
        glRenderbufferStorageEXT (GL_RENDERBUFFER_EXT, GL_DEPTH_COMPONENT, w, h)
        
        # Create texture to render to
        debug('allocating dynamic texture')
        texture = glGenTextures (1)
        glBindTexture (GL_TEXTURE_2D, texture)
        glTexParameteri (GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri (GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        # Question: why do we need to upload a texture that we are rendering *to* ?
        glTexImage2D (GL_TEXTURE_2D, 0, GL_RGB, w, h, 0,
                        GL_RGB, GL_UNSIGNED_BYTE, bits)

        # store variables for later use.
        self.texture = texture
        self.framebuffer = framebuffer
        self.depthbuffer = depthbuffer
        self.image = image
        self.bits = bits
        self.size = size
        self.texture_size = (w,h)
        self.upside_down = upside_down

        # set the screen as output
        glBindRenderbufferEXT (GL_RENDERBUFFER_EXT, 0)
        glBindFramebufferEXT (GL_FRAMEBUFFER_EXT, 0)

    def __del__(self):
        glBindRenderbufferEXT (GL_RENDERBUFFER_EXT, 0)
        glBindFramebufferEXT (GL_FRAMEBUFFER_EXT, 0)
        glDeleteTextures (1,[self.texture])
        glDeleteFramebuffersEXT (1, [self.framebuffer])
        glDeleteRenderbuffersExt (1, [self.depthbuffer])

    def setup(self):
        glBindRenderbufferEXT (GL_RENDERBUFFER_EXT, self.depthbuffer)
        glBindFramebufferEXT(GL_FRAMEBUFFER_EXT,  self.framebuffer)
        glFramebufferTexture2DEXT (GL_FRAMEBUFFER_EXT, GL_COLOR_ATTACHMENT0_EXT,
                                   GL_TEXTURE_2D,  self.texture, 0);
        glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT, GL_DEPTH_ATTACHMENT_EXT, 
                                     GL_RENDERBUFFER_EXT,  self.depthbuffer);
                                    
        status = glCheckFramebufferStatusEXT (GL_FRAMEBUFFER_EXT);
        if status != GL_FRAMEBUFFER_COMPLETE_EXT:
            raise RuntimeError( "Error in framebuffer activation")

        # Re-orient viewport
        glViewport (0, 0, self.size[0], self.size[1])
        glMatrixMode (GL_PROJECTION)
        glLoadIdentity()
        gluPerspective (40.,self.size[0]/float(self.size[1]),1.,40.)
        glMatrixMode (GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt (0,0,10, 0,0,0, 0,1,0)

    def render(self):
        # TODO: Can we get away with glFlush?
        glFinish() #renders to our framebuffer

        # read back the framebuffer to self.image
        glBindTexture (GL_TEXTURE_2D, self.texture)
        w,h = self.texture_size
        data = glReadPixels (0, 0, w, h, GL_RGB,  GL_UNSIGNED_BYTE)
        rval = numpy.fromstring(data, dtype='uint8', count=w*h*3).reshape((w,h,3))
        if self.size != self.texture_size:
            rval = rval[:self.size[0], :self.size[1],:]

        # return to default state of screen rendering
        glBindRenderbufferEXT (GL_RENDERBUFFER_EXT, 0)
        glBindFramebufferEXT(GL_FRAMEBUFFER_EXT,  0)
        if self.upside_down:
            return rval
        else:
            return rval[::-1,:,:]

class OpenGlMovieFromImage(theano.Op):
    """Helper base class to factor code used by Ops that want to make a movie from an input
    image, using OpenGL.  The subclass specifies how to actually make the movie.
    """

    def __init__(self, width, height, upside_down=False):
        """To set up the renderer, we need to know the frame size of the images.
        Setting up the renderer for each image is much slower.
        """
        init_GL() #global initialization is no-op after first call
        self.width=width
        self.height=height
        self.upside_down=upside_down

        self.renderer = None
        # Delay construction of renderer until after merge-optimization
        #PBufRenderer((width, height), upside_down=upside_down)

        #TODO: put texture into output state as reusable resource
        self.texture = glGenTextures(1)

    def __del__(self):
        glDeleteTextures (1,[self.texture])

    def __eq__(self, other):
        return type(self) == type(other) \
                and self.width == other.width \
                and self.height == other.height \
                and self.upside_down == other.upside_down

    def __hash__(self):
        return hash(type(self)) ^ hash(self.width) ^ hash(self.height) ^ hash(self.upside_down)

    def make_node(self, x, istate):
        _x = theano.tensor.as_tensor_variable(x)
        if _x.type.dtype != 'uint8':
            raise TypeError('must be 2- or 3-tensor of uint8', x)
        if _x.type.broadcastable != (False, False) \
                and _x.type.broadcastable != (False, False, False):
            raise TypeError('must be a 2- or 3-tensor of uint8', x)
        if not isinstance(istate, theano.Variable):
            raise TypeError("variable expected", istate)
        o_type = theano.tensor.TensorType(dtype='uint8', broadcastable=[False, False, False, False])
        state_type = theano.gof.type.generic
        return theano.Apply(self, [x, istate], [o_type(), state_type()])

    def perform(self, node, (x, istate), (z_storage, ostate_storage)):
        if self.renderer is None:
            self.renderer = PBufRenderer((self.width, self.height), upside_down=self.upside_down)

        ostate = copy.deepcopy(istate)
        self.renderer.setup() 

        glBindTexture(GL_TEXTURE_2D, self.texture)   # 2d texture (x and y size)
        load_texture(x)

        z = numpy.zeros(self.z_shape, dtype='uint8')
        for i in xrange(self.n_frames):
            self.perform_set_state(istate, ostate, i)
            self.perform_display(x, ostate, i)
            di = self.renderer.render()
            z[i] = di

        # store output images 
        z_storage[0] = z

        # store next state ostate_storage
        ostate_storage[0] = ostate

class ImageOnSpinningCube(OpenGlMovieFromImage):
    def __init__(self, (n_frames, width, height), upside_down=False):
        super(ImageOnSpinningCube, self).__init__(width, height, upside_down=upside_down)
        self.n_frames = n_frames
        self.z_shape = (n_frames, width, height, 3)

    def __eq__(self, other):
        return super(ImageOnSpinningCube, self).__eq__(other) \
                and self.n_frames == other.n_frames \

    def __hash__(self):
        return super(ImageOnSpinningCube, self).__hash__() ^ hash(self.n_frames)

    def new_state(self, speed=10):
        return dict(
                rot=numpy.asarray((0.,0.,0.)),
                drot=numpy.asarray((speed,speed,speed)),
                )

    def perform_set_state(self, istate, ostate, iter):
        ostate['rot'] = istate['rot'] + istate['drot'] * iter

    def perform_display(self, x, ostate, i):
        # retrieve some state variables related to rendering
        xrot,yrot,zrot = ostate['rot']
        dxrot,dyrot,dzrot = ostate['drot']

        # load x as a texture
        glBindTexture(GL_TEXTURE_2D, self.texture)   # 2d texture (x and y size)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL)

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)      # Clear The Screen And The Depth Buffer
        glLoadIdentity()                                        # Reset The View
        glTranslatef(0.0,0.0,-5.0)                      # Move Into The Screen

        glRotatef(xrot,1.0,0.0,0.0)                     # Rotate The Cube On It's X Axis
        glRotatef(yrot,0.0,1.0,0.0)                     # Rotate The Cube On It's Y Axis
        glRotatef(zrot,0.0,0.0,1.0)                     # Rotate The Cube On It's Z Axis

        glBegin(GL_QUADS)                           # Start Drawing The Cube
        
        # Front Face (note that the texture's corners have to match the quad's corners)
        glTexCoord2f(0.0, 0.0); glVertex3f(-1.0, -1.0,  1.0)    # Bottom Left Of The Texture and Quad
        glTexCoord2f(1.0, 0.0); glVertex3f( 1.0, -1.0,  1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f( 1.0,  1.0,  1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f(-1.0,  1.0,  1.0)    # Top Left Of The Texture and Quad
        
        # Back Face
        glTexCoord2f(1.0, 0.0); glVertex3f(-1.0, -1.0, -1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f(-1.0,  1.0, -1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f( 1.0,  1.0, -1.0)    # Top Left Of The Texture and Quad
        glTexCoord2f(0.0, 0.0); glVertex3f( 1.0, -1.0, -1.0)    # Bottom Left Of The Texture and Quad
        
        # Top Face
        glTexCoord2f(0.0, 1.0); glVertex3f(-1.0,  1.0, -1.0)    # Top Left Of The Texture and Quad
        glTexCoord2f(0.0, 0.0); glVertex3f(-1.0,  1.0,  1.0)    # Bottom Left Of The Texture and Quad
        glTexCoord2f(1.0, 0.0); glVertex3f( 1.0,  1.0,  1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f( 1.0,  1.0, -1.0)    # Top Right Of The Texture and Quad
        
        # Bottom Face       
        glTexCoord2f(1.0, 1.0); glVertex3f(-1.0, -1.0, -1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f( 1.0, -1.0, -1.0)    # Top Left Of The Texture and Quad
        glTexCoord2f(0.0, 0.0); glVertex3f( 1.0, -1.0,  1.0)    # Bottom Left Of The Texture and Quad
        glTexCoord2f(1.0, 0.0); glVertex3f(-1.0, -1.0,  1.0)    # Bottom Right Of The Texture and Quad
        
        # Right face
        glTexCoord2f(1.0, 0.0); glVertex3f( 1.0, -1.0, -1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f( 1.0,  1.0, -1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f( 1.0,  1.0,  1.0)    # Top Left Of The Texture and Quad
        glTexCoord2f(0.0, 0.0); glVertex3f( 1.0, -1.0,  1.0)    # Bottom Left Of The Texture and Quad
        
        # Left Face
        glTexCoord2f(0.0, 0.0); glVertex3f(-1.0, -1.0, -1.0)    # Bottom Left Of The Texture and Quad
        glTexCoord2f(1.0, 0.0); glVertex3f(-1.0, -1.0,  1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f(-1.0,  1.0,  1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f(-1.0,  1.0, -1.0)    # Top Left Of The Texture and Quad
        
        glEnd();                                # Done Drawing The Cube

def image_on_spinning_cube(x, shape, upside_down=False):
    op = ImageOnSpinningCube(shape, upside_down=upside_down)
    istate = shared(op.new_state())
    z, ostate = op(x, istate)
    return z, {istate: ostate}

class BrownianCamera(OpenGlMovieFromImage):
    def __init__(self, (n_frames, width, height), upside_down=False):
        super(BrownianCamera, self).__init__(width, height, upside_down=upside_down)
        self.n_frames = n_frames
        self.z_shape = (n_frames, width, height, 3)

    def __eq__(self, other):
        return super(self.__class__, self).__eq__(other) \
                and self.n_frames == other.n_frames \

    def __hash__(self):
        return super(self.__class__, self).__hash__() ^ hash(self.n_frames)

    def new_state(self, pos_jitter=(.01,.01,.03), rot_jitter=(4.,4.,4.), seed=23424):
        return dict(
                pos_jitter=numpy.asarray(pos_jitter),
                rot_jitter=numpy.asarray(rot_jitter),
                pos0=numpy.asarray((0.,0.,-4.0)),
                rot0=numpy.asarray((0.,0.,0.)),
                alpha=0.1,
                # dynamic things
                pos=numpy.asarray((0.,0.,-4.0)),
                dpos=numpy.asarray((0.,0.,0.)),
                ddpos=numpy.asarray((0.,0.,0.)),
                rot=numpy.asarray((0.,0.,0.)),
                drot=numpy.asarray((0.,0.,0.)),
                ddrot=numpy.asarray((0.,0.,0.)),
                rng = numpy.random.RandomState(seed),
                )

    def perform_set_state(self, istate, ostate, iter):
        alpha = ostate['alpha']
        if iter == 0:
            ostate['pos'] = ostate['pos0'].copy()
            ostate['dpos'] *= 0
            ostate['rot'] = ostate['rot0'].copy()
            ostate['drot'] *= 0
        ostate['ddpos'] = ostate['rng'].uniform(low=-1,high=1,size=3) * ostate['pos_jitter']
        ostate['ddrot'] = ostate['rng'].uniform(low=-1,high=1,size=3) * ostate['rot_jitter']
        ostate['dpos'] += ostate['ddpos']
        ostate['drot'] += ostate['ddrot']
        ostate['pos'] = (1-alpha)*(ostate['pos'] + ostate['dpos']) + alpha * ostate['pos0']
        ostate['rot'] = (1-alpha)*(ostate['rot'] + ostate['drot']) + alpha * ostate['rot0']

    def perform_display(self, x, ostate, i):
        # retrieve some state variables related to rendering
        xrot,yrot,zrot = ostate['rot']
        xpos,ypos,zpos = ostate['pos']

        # load x as a texture
        glBindTexture(GL_TEXTURE_2D, self.texture)   # 2d texture (x and y size)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL)

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)      # Clear The Screen And The Depth Buffer
        glLoadIdentity()                                        # Reset The View
        glTranslatef(xpos,ypos,zpos)                      # Move Into The Screen

        glRotatef(xrot,1.0,0.0,0.0)                     # Rotate The Cube On It's X Axis
        glRotatef(yrot,0.0,1.0,0.0)                     # Rotate The Cube On It's Y Axis
        glRotatef(zrot,0.0,0.0,1.0)                     # Rotate The Cube On It's Z Axis

        glBegin(GL_QUADS)                           # Start Drawing The Cube
        
        # Front Face (note that the texture's corners have to match the quad's corners)
        glTexCoord2f(0.0, 0.0); glVertex3f(-1.0, -1.0,  1.0)    # Bottom Left Of The Texture and Quad
        glTexCoord2f(1.0, 0.0); glVertex3f( 1.0, -1.0,  1.0)    # Bottom Right Of The Texture and Quad
        glTexCoord2f(1.0, 1.0); glVertex3f( 1.0,  1.0,  1.0)    # Top Right Of The Texture and Quad
        glTexCoord2f(0.0, 1.0); glVertex3f(-1.0,  1.0,  1.0)    # Top Left Of The Texture and Quad
        
        glEnd();                                # Done Drawing The Cube

_brownian_camera_ops = {}
def brownian_camera(x, shape, upside_down=False, seed=8234, speed=1.0):
    if (shape, upside_down) not in _brownian_camera_ops:
        _brownian_camera_ops[(shape, upside_down)] = BrownianCamera(shape, upside_down=upside_down)
    op = _brownian_camera_ops[(shape, upside_down)]
    istate = shared(op.new_state(seed=seed))
    istate.value['pos_jitter'] *= speed
    istate.value['rot_jitter'] *= speed
    z, ostate = op(x, istate)
    return z, [(istate, ostate)]


def _dump_to_file(fn, filename='out.pkl', nexamples=1000, n_frames=10, **kwargs):
    logging.basicConfig(level=logging.INFO, stream=sys.stderr)
    import cPickle, time

    from MNIST import mnist
    i = theano.tensor.iscalar()
    z, z_updates = fn(mnist(i%50000, 'train', rasterized=False, dtype='uint8')[0], (n_frames, 28,28), **kwargs)
    f = function([i], z[:,:,:,0], updates=z_updates)

    t0 = time.time()
    rval = []
    for j in xrange(nexamples):
        if 0 == j  % 100: print >> sys.stderr, j
        rval.append(f(j))
    dt = time.time() - t0
    info('Generating ', nexamples, 'examples took', dt, 'seconds.')
    info('Generation rate:', nexamples/dt, 'examples per second.')
    info('Generated ', nexamples*n_frames, 'frames')
    info('Generation rate:', nexamples*n_frames/dt, 'frames per second.')

    cPickle.dump(rval, file(filename, 'w'), protocol=cPickle.HIGHEST_PROTOCOL)
def spinning_cube_dump(filename='spinning_cube.pkl', *args, **kwargs):
    return _dump_to_file(fn=image_on_spinning_cube, filename=filename, *args, **kwargs)
def brownian_camera_dump(filename='brownian_camera.pkl', *args, **kwargs):
    return _dump_to_file(fn=brownian_camera, filename=filename, *args, **kwargs)
def brownian_camera_dumpN(filename='brownian_cameraN.pkl', nexamples=10, n_frames=5,
        n_movies=10, img_shape=(28,28), **kwargs):
    logging.basicConfig(level=logging.INFO, stream=sys.stderr)
    import cPickle, time
    from MNIST import mnist

    s_idx = theano.tensor.iscalar()
    inputs_updates = [brownian_camera(
            x=mnist(s_idx*n_movies+i, 'train', rasterized=False, dtype='uint8')[0],
            shape=(n_frames,)+img_shape,
            seed=234234+i, **kwargs)
        for i in xrange(n_movies)]
    s_input = theano.tensor.stack(*(input for (input,update) in inputs_updates))\
            .reshape((n_movies*n_frames,)+img_shape+(3,))
    s_updates = []
    for i,u in inputs_updates:
        s_updates.extend(u)
    print s_updates
    f = function([s_idx], s_input, updates=s_updates)

    t0 = time.time()
    rval = []
    for j in xrange(nexamples):
        if 0 == j  % 1000: print >> sys.stderr, j
        rval.append(f(j))
    dt = time.time() - t0
    info('Generating ', nexamples, 'examples took', dt, 'seconds.')
    info('Generation rate:', nexamples/dt, 'examples per second.')
    info('Generated ', nexamples*n_movies*n_frames, 'frames')
    info('Generation rate:', nexamples*n_movies*n_frames/dt, 'frames per second.')

    cPickle.dump(rval, file(filename, 'w'))


def glviewer_from_file(filename='out.pkl'):
    logging.basicConfig(level=logging.DEBUG, stream=sys.stderr)
    import cPickle
    rval = cPickle.load(file(filename))
    from glviewer import GlViewer
    GlViewer(rval.__getitem__).main()