#!/usr/bin/env python

import sys
import math
import thread
from PyOSG import *

class MyUpdateCallback(osg.NodeCallback):
    def __init__(self, subgraph):
        osg.NodeCallback.__init__(self)
        self._subgraph = subgraph

    def apply(self, node, nv):
        # traverse the subgraph to update any nodes
        if self._subgraph:
            self._subgraph.accept(nv)

        # must traverse the Node's subgraph
        self.traverse(node,nv)

class MyCullCallback(osg.NodeCallback):
    def __init__(self, subgraph, attr):
        osg.NodeCallback.__init__(self)
        self._subgraph = subgraph
        self._texture = None
        self._image = None
        self._localState = osg.StateSet()
        if type(attr) == type(osg.Texture2D()):
            self._texture = attr
        if type(attr) == type(osg.Image()):
            self._image = attr

    def apply(self, node, nv):
        cullVisitor = osgUtil.asCullVisitor(nv)
        if cullVisitor and (self._texture or self._image) and self._subgraph:
            self.doPreRender(node, cullVisitor)

            # must traverse the subgraph
            self.traverse(node,nv)
        else:
            # must traverse the subgraph
            self.traverse(node,nv)

    def doPreRender(self, node, cv):
        bs = self._subgraph.getBound()
        if not bs:
            print "bb invalid ", self._subgraph
            return

        # create the render to texture stage.
        rtts = osgUtil.RenderToTextureStage()

        # set up lighting.
        # currently ignore lights in the scene graph itself..
        # will do later.
        previous_stage = cv.getCurrentRenderBin().getStage()

        # set up the background color and clear mask.
        rtts.setClearColor(osg.Vec4(0.1,0.1,0.3,1.0))
        rtts.setClearMask(previous_stage.getClearMask())

        # set up to charge the same RenderStageLighting is the parent previous stage.
        rtts.setRenderStageLighting(previous_stage.getRenderStageLighting())


        # record the render bin, to be restored after creation
        # of the render to text
        previousRenderBin = cv.getCurrentRenderBin()

        # set the current renderbin to be the newly created stage.
        cv.setCurrentRenderBin(rtts)


        znear = 1.0*bs.radius()
        zfar  = 3.0*bs.radius()

        # 2:1 aspect ratio as per flag geomtry below.
        top   = 0.25*znear
        right = 0.5*znear

        znear *= 0.9
        zfar *= 1.1

        # set up projection.
        projection = osg.RefMatrix()
        projection.makeFrustum(-right,right,-top,top,znear,zfar)

        cv.pushProjectionMatrix(projection)

        matrix = osg.RefMatrix()
        matrix.makeLookAt(bs.center()+osg.Vec3(0.0,2.0,0.0)*bs.radius(),bs.center(),osg.Vec3(0.0,0.0,1.0))

        cv.pushModelViewMatrix(matrix)

        cv.pushStateSet(self._localState)
        # traverse the subgraph
        self._subgraph.accept(cv)

        cv.popStateSet()

        # restore the previous model view matrix.
        cv.popModelViewMatrix()

        # restore the previous model view matrix.
        cv.popProjectionMatrix()

        # restore the previous renderbin.
        cv.setCurrentRenderBin(previousRenderBin)

        if rtts.getRenderGraphList().size()==0 and rtts.getRenderBinList().size()==0:
            # getting to this point means that all the subgraph has been
            # culled by small feature culling or is beyond LOD ranges.
            return

        height = 256
        width  = 512


        viewport = cv.getViewport()

        # offset the impostor viewport from the center of the main window
        # viewport as often the edges of the viewport might be obscured by
        # other windows, which can cause image/reading writing problems.
        center_x = viewport.x()+viewport.width()/2
        center_y = viewport.y()+viewport.height()/2

        new_viewport = osg.Viewport()
        new_viewport.setViewport(center_x-width/2,center_y-height/2,width,height)
        rtts.setViewport(new_viewport)

        self._localState.setAttribute(new_viewport)

        # and the render to texture stage to the current stages
        # dependancy list.
        cv.getCurrentRenderBin().getStage().addToDependencyList(rtts)

        # if one exist attach texture to the RenderToTextureStage.
        if self._texture:
            rtts.setTexture(self._texture)

        # if one exist attach image to the RenderToTextureStage.
        if self._image:
            rtts.setImage(self._image)

class MyGeometryCallback(osg.Drawable.UpdateCallback, osg.Drawable.AttributeFunctor):
    """ call back which creates a deformation field to oscilate the model."""
    def __init__(self, o, x, y, z, period, xphase, amplitude):
        osg.Drawable.UpdateCallback.__init__(self)
        osg.Drawable.AttributeFunctor.__init__(self)

        self._firstCall = True
        self._startTime = 0.0
        self._time = 0.0
        self._period = period
        self._xphase = xphase
        self._amplitude = amplitude
        self._origin = o
        self._xAxis = x
        self._yAxis = y
        self._zAxis = z

    def update (self, nv, drawable):
        fs = nv.getFrameStamp()
        referenceTime = fs.getReferenceTime()
        if self._firstCall:
            self._firstCall = False
            self._startTime = referenceTime

        self._time = referenceTime - self._startTime

        drawable.accept(self)
        drawable.dirtyBound()

        geometry = drawable.asGeometry()
        if geometry:
            osgUtil.SmoothingVisitor.smooth(geometry)

    def apply(self, drawtype, count, begin):
        if drawtype == osg.Drawable.VERTICES:
            TwoPI = 2.0 * osg.PI
            phase = -self._time/self._period

            for itr in begin[:count]:
                dv = itr - self._origin
                local = osg.Vec3(dv*self._xAxis,dv*self._yAxis,dv*self._zAxis)
                local[2] = local._x*self._amplitude * math.sin(TwoPI*(phase+local._x*self._xphase));

                # Don't assign to itr, since this will create a new object
                itr.set(self._origin + self._xAxis*local._x+ self._yAxis*local._y+ self._zAxis*local._z)

def createPreRenderSubGraph(subgraph):
    if not subgraph:
        return 0

    # create the quad to visualize.
    polyGeom = osg.Geometry()

    polyGeom.setSupportsDisplayList(False)

    origin = osg.Vec3(0.0,0.0,0.0)
    xAxis = osg.Vec3(1.0,0.0,0.0)
    yAxis = osg.Vec3(0.0,0.0,1.0)
    zAxis = osg.Vec3(0.0,-1.0,0.0)
    height = 100.0
    width = 200.0
    noSteps = 20

    vertices = osg.Vec3Array()
    bottom = osg.Vec3()
    bottom.set(origin)

    top = osg.Vec3()
    top.set(origin)

    top[2] += height
    dv = xAxis * (width/(noSteps-1.0))

    texcoords = osg.Vec2Array()
    bottom_texcoord = osg.Vec2(0.0,0.0)
    top_texcoord = osg.Vec2(0.0,1.0)
    dv_texcoord = osg.Vec2(1.0/(float)(noSteps-1),0.0)

    for i in range(noSteps):
        vertices.push_back(top)
        vertices.push_back(bottom)
        top+=dv
        bottom+=dv

        texcoords.push_back(top_texcoord)
        texcoords.push_back(bottom_texcoord)
        top_texcoord+=dv_texcoord
        bottom_texcoord+=dv_texcoord

    # pass the created vertex array to the points geometry object.
    polyGeom.setVertexArray(vertices)

    polyGeom.setTexCoordArray(0,texcoords)


    colors = osg.Vec4Array()
    colors.push_back(osg.Vec4(1.0,1.0,1.0,1.0))
    polyGeom.setColorArray(colors)
    polyGeom.setColorBinding(osg.Geometry.BIND_OVERALL)

    polyGeom.addPrimitiveSet(osg.DrawArrays(osg.PrimitiveSet.QUAD_STRIP,0,vertices.size()))

    # new we need to add the texture to the Drawable, we do so by creating a
    # StateSet to contain the Texture StateAttribute.
    stateset = osg.StateSet()

    texture = osg.Texture2D()
    texture.setFilter(osg.Texture2D.MIN_FILTER,osg.Texture2D.LINEAR)
    texture.setFilter(osg.Texture2D.MAG_FILTER,osg.Texture2D.LINEAR)
    stateset.setTextureAttributeAndModes(0, texture,osg.StateAttribute.ON)

    polyGeom.setStateSet(stateset)

    polyGeom.setUpdateCallback(MyGeometryCallback(origin,xAxis,yAxis,zAxis,1.0,1.0/width,0.2))

    geode = osg.Geode()
    geode.addDrawable(polyGeom)

    parent = osg.Group()

    parent.setUpdateCallback(MyUpdateCallback(subgraph))

    parent.setCullCallback(MyCullCallback(subgraph,texture))

    parent.addChild(geode)

    return parent

def main(argv):
    # use an ArgumentParser object to manage the program arguments.
    arguments =  osg.ArgumentParser(argv)

    # set up the usage document, in case we need to print out how to use this program.
    arguments.getApplicationUsage().setApplicationName(arguments.getApplicationName())
    arguments.getApplicationUsage().setDescription(arguments.getApplicationName()+" is the example which demonstrates pre rendering of scene to a texture, and then apply this texture to geometry.")
    arguments.getApplicationUsage().setCommandLineUsage(arguments.getApplicationName()+" [options] filename ...")

    # construct the viewer.
    viewer = osgProducer.Viewer(arguments)

    # set up the value with sensible default event handlers.
    viewer.setUpViewer(osgProducer.Viewer.STANDARD_SETTINGS)

    # get details on keyboard and mouse bindings used by the viewer.
    viewer.getUsage(arguments.getApplicationUsage())

    # if user request help write it out to cout.
    if arguments.read("-h") or arguments.read("--help"):
        arguments.getApplicationUsage().write()
        return 1

    # any option left unread are converted into errors to write out later.
    arguments.reportRemainingOptionsAsUnrecognized()

    # report any errors if they have occured when parsing the program aguments.
    if arguments.errors():
        arguments.writeErrorMessages()
        return 1

    if arguments.argc()<=1:
        arguments.getApplicationUsage().write(osg.ApplicationUsage.COMMAND_LINE_OPTION)
        return 1

    # read the scene from the list of file specified commandline args.
    loadedModel = osgDB.readNodeFiles(arguments)

    # if no model has been successfully loaded report failure.
    if not loadedModel: 
        print arguments.getApplicationName(), " No data loaded"
        return 1


    # create a transform to spin the model.
    loadedModelTransform = osg.MatrixTransform()
    loadedModelTransform.addChild(loadedModel)

    nc = osgUtil.TransformCallback(loadedModelTransform.getBound().center(),osg.Vec3(0.0,0.0,1.0),osg.inDegrees(45.0))
    loadedModelTransform.setUpdateCallback(nc)

    rootNode = osg.Group()
    # rootNode.addChild(loadedModelTransform)
    rootNode.addChild(createPreRenderSubGraph(loadedModelTransform))

    # set the scene to render
    viewer.setSceneData(rootNode)

    # create the windows and run the threads.
    viewer.realize()
    # viewer.realize(Producer.CameraGroup.SingleThreaded)

    while not viewer.done():
        # wait for all cull and draw threads to complete.
        viewer.sync()

        # update the scene by traversing it with the the update visitor which will
        # call all node update callbacks and animations.
        viewer.update()

        # fire off the cull and draw traversals of the scene.
        viewer.frame()
    # wait for all cull and draw threads to complete before exit.
    viewer.sync()

    return 0

if __name__ == "__main__":
    main(sys.argv)