source: trunk/anuga_core/anuga/visualiser/visualiser.py @ 9679

from threading import Thread
from Queue import Queue
from Tkinter import Tk, Button, Frame, N, E, S, W
from types import FunctionType, TupleType
from vtk import vtkActor, vtkCubeAxesActor2D, vtkDelaunay2D, vtkFloatArray, vtkPoints, vtkPolyData, vtkPolyDataMapper, vtkRenderer
from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget

class Visualiser(Thread):
    """Superclass of both the realtime and offline VTK visualisers
    """
    def __init__(self, source):
        Thread.__init__(self)

        self.source = source

        # Structures for Height Based quantities
        self.height_quantities = []
        self.height_zScales = {}
        self.height_dynamic = {}
        self.height_offset = {}
        self.height_opacity = {}
        self.height_wireframe = {}

        # Structures for colouring quantities
        self.colours_height = {}

        # Structures used for VTK
        self.vtk_actors = {}
        self.vtk_axesSet = False
        self.vtk_drawAxes = False
        self.vtk_mappers = {}
        self.vtk_polyData = {}

        # A list of operations to be performed on the cube axes. Type: [(func, (args))]
        self.conf_axesAlterations = []
        # A list of all polygons to overlay. Type: [([coords], height, (colour)]
        self.conf_overlaidPolygons = []
        # A list of alterations to be performed on the Tk root. Type: [(func, (args))]
        self.conf_tkAlterations = []

    def run(self):
        self.vtk_renderer = vtkRenderer()
        self.setup_gui()
        self.setup_grid()

        # Handle any deferred configuration
        # Overlaid polygons
        for args in self.conf_overlaidPolygons:
            self.overlay_polygon_internal(*args)
        # Draw (and maybe alter) the axes
        if self.vtk_drawAxes:
            self.vtk_axes = vtkCubeAxesActor2D()
            # Perform all of the alterations required, by applying func to the vtk_axes instance (with the given args).
            for func, args in self.conf_axesAlterations:
                func(*((self.vtk_axes,) + args))
        # Alter the Tk root as necessary.
        for func, args in self.conf_tkAlterations:
            func(*((self.tk_root,) + args))
        # Finished with deferred configuration.

        # Draw Height Quantities
        for q in self.height_quantities:
            self.update_height_quantity(q, self.height_dynamic[q])
            self.draw_height_quantity(q)
            
        self.tk_root.mainloop()

    def redraw_quantities(self):
        """Redraw all dynamic quantities.
        """
        # Height quantities
        for q in self.height_quantities:
            if (self.height_dynamic[q]):
                self.update_height_quantity(q, self.height_dynamic[q])
                self.draw_height_quantity(q)
        if self.vtk_drawAxes is True:
            self.draw_axes()

    # --- Axes --- #
        
    def render_axes(self):
        """Intstruct the visualiser to render cube axes around the render.
        """
        self.vtk_drawAxes = True

    def draw_axes(self):
        """Update the 3D bounds on the axes and add them to the pipeline if not yet connected.
        """
        self.vtk_axes.SetBounds(self.get_3d_bounds())
        if not self.vtk_axesSet:
            self.vtk_axesSet = True
            self.vtk_axes.SetCamera(self.vtk_renderer.GetActiveCamera())
            self.vtk_renderer.AddActor(self.vtk_axes)
            self.vtk_renderer.ResetCamera(self.get_3d_bounds())
        
    def alter_axes(self, func, args):
        """Attempt to apply the function 'func' with args tuple 'args' to the
        vtkCubeAxesActor2D instance set up by render_axes. This is done this way to ensure
        the axes setup is handled in the visualiser thread.

        Example call:
        from vtk import vtkCubeAxesActor2D
        alter_axes(vtkCubeAxesActor2D.SetNumberOfPoints, (5,))
        """
        self.conf_axesAlterations.append((func, args))
            
    # --- Height Based Rendering --- #

    def setup_grid(self):
        """Create the vtkCellArray instance that represents the
        triangles. Subclasses are expected to override this function
        to read from their source as appropriate. The vtkCellArray should
        be stored to self.vtk_cells.
        """
        pass

    def render_quantity_height(self, quantityName, zScale=1.0, offset=0.0, opacity=1.0, dynamic=True, wireframe=False):
        """Instruct the visualiser to render a quantity using the
        value at a point as its height.  The value at each point is
        multiplied by z_scale and is added to offset, and if
        dynamic=False, the quantity is not recalculated on each
        update.
        """
        self.height_quantities.append(quantityName)
        self.height_zScales[quantityName] = zScale
        self.height_offset[quantityName] = offset
        self.height_dynamic[quantityName] = dynamic
        self.height_opacity[quantityName] = opacity
        self.height_wireframe[quantityName] = wireframe

    def update_height_quantity(self, quantityName, dynamic=True):
        """Create a vtkPolyData object and store it in
        self.vtk_polyData[quantityName]. Subclasses are expected to override this
        function.
        """
        pass

    def get_3d_bounds(self):
        """Get the minimum and maximum bounds for the x, y and z directions.
        Return as a list of double in the order (xmin, xmax, ymin, ymax, zmin, zmax),
        suitable for passing to vtkCubeAxesActor2D::SetRanges(). Subclasses are expected
        to override this function.
        """
        pass



    def store_height_quantity(self, quantityName, fileName=None):

        if fileName == None:
            fileName = quantityName + '.vtk'

        quantity_polyData = self.vtk_polyData[quantityName]

        import vtk
        w = vtk.vtkPolyDataWriter()
        #print quantity_polyData
        w.SetInput(quantity_polyData)
        w.SetFileName(fileName)
        w.Write()



    def draw_height_quantity(self, quantityName):
        """Use the vtkPolyData and prepare/update the rest of the VTK
        rendering pipeline.
        """
        if self.vtk_mappers.has_key(quantityName):
            mapper = self.vtk_mappers[quantityName]
        else:
            mapper = self.vtk_mappers[quantityName] = vtkPolyDataMapper()
        mapper.SetInput(self.vtk_polyData[quantityName])
        mapper.Update()

        if not self.vtk_actors.has_key(quantityName):
            actor = self.vtk_actors[quantityName] = vtkActor()
            actor.GetProperty().SetOpacity(self.height_opacity[quantityName])
            if self.height_wireframe[quantityName]:
                actor.GetProperty().SetRepresentationToWireframe()
            actor.SetMapper(mapper)
            self.vtk_renderer.AddActor(actor)
        else:
            actor = self.vtk_actors[quantityName]

        if self.colours_height.has_key(quantityName):
            colour = self.colours_height[quantityName]
            if type(colour) == TupleType:
                if type(colour[0]) == FunctionType:
                    # It's a function, so take colour[1] as the
                    # lower bound on the scalar range and
                    # colour[2] as the upper bound on the scalar
                    # range.
                    scalars = vtkFloatArray()

                    map(scalars.InsertNextValue, colour[0](self.build_quantity_dict()))
                    self.vtk_polyData[quantityName].GetPointData().SetScalars(scalars)
                    mapper.SetScalarRange(colour[1:])
                    mapper.Update()
                else:
                    # It's a 3-tuple representing an RGB value.
                    actor.GetProperty().SetColor(colour)
            else:
                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
        else:
            actor.GetProperty().SetColor(0.5, 0.5, 0.5)

    # --- Colour Coding --- #

    def build_quantity_dict(self):
        """Build and return a dictionary mapping quantity name->Numeric array of vertex
        values for that quantity. Subclasses are expected to override
        this function."""
        pass

    def colour_height_quantity(self, quantityName, colour=(0.5, 0.5, 0.5)):
        """Add colouring to a height based quantity.

        The colour parameter can be one of the following:
        - a 3-tuple of values in [0,1] to specify R, G, B values
        - a 3-tuple of values:
          - a function that takes a dictionary mapping quantity name->Numeric array of vertex values.
            This function returns a list of vertex values to be used in the colour coding.
          - a float for the lower bound on the colouring
          - a float for the upper bound on the colouring
        """
        self.colours_height[quantityName] = colour

    # --- Overlaid Polygons --- #

    def overlay_polygon(self, coords, height=0.0, colour=(1.0, 0.0, 0.0)):
        """Add a polygon to the output of the visualiser.

        coords is a list of 2-tuples representing x and y coordinates.
        These are triangulated by vtkDelaunay2D.

        height is the z-value given to all points.

        colour is the colour of the polygon, as a 3-tuple representing
        r, g, b values between 0 and 1."""
        self.conf_overlaidPolygons.append((coords, height, colour))

    def overlay_polygon_internal(self, coords, height, colour):
        """Add a polygon to the output of the visualiser.

        coords is a list of 2-tuples representing x and y coordinates.
        These are triangulated by vtkDelaunay2D.

        height is the z-value given to all points.

        colour is the colour of the polygon, as a 3-tuple representing
        r, g, b values between 0 and 1.

        This function should not be called from outside the visualiser thread.
        Use overlay_polygon instead.
    
        """
        points = vtkPoints()
        for coord in coords:
            points.InsertNextPoint(coord[0], coord[1], height)
        profile = vtkPolyData()
        profile.SetPoints(points)
        delny = vtkDelaunay2D()
        delny.SetInput(profile)
        mesh = vtkPolyDataMapper()
        mesh.SetInput(delny.GetOutput())
        actor = vtkActor()
        actor.SetMapper(mesh)
        actor.GetProperty().SetColor(colour)
        self.vtk_renderer.AddActor(actor)
        
    # --- Vector Fields --- #

    # --- GUI Setup --- #

    def setup_gui(self):
        self.tk_root = Tk()
        self.tk_root.title("Visualisation")
        self.tk_root.after(100, self.redraw)
        self.tk_root.bind("<Destroy>", self.destroyed)
        self.tk_root.grid_rowconfigure(0, weight=1)
        self.tk_root.grid_columnconfigure(0, weight=1)

        self.tk_renderWidget = vtkTkRenderWidget(self.tk_root, width=400, height=400)
        self.tk_renderWidget.grid(row=0, column=0, sticky=N+S+E+W)
        self.tk_controlFrame = Frame(self.tk_root)
        self.tk_controlFrame.grid(row=1, column=0, sticky=E+W)
        self.tk_controlFrame.grid_rowconfigure(0, weight=1)
        self.tk_controlFrame.grid_columnconfigure(0, weight=1)
        
        self.tk_quit = Button(self.tk_controlFrame, text="Quit", command=self.shutdown)
        self.tk_quit.grid(row=0, column=0, sticky=E+W)
        self.tk_renderWidget.GetRenderWindow().AddRenderer(self.vtk_renderer)

    def alter_tkroot(self, func, args):
        """Apply func, with arguments tuple args to the root tk window for this visualiser.
        """
        self.conf_tkAlterations.append((func, args))

    # --- GUI Events --- #

    def destroyed(self, event):
        if event.widget == self.tk_root:
            self.shutdown()

    def redraw(self):
        self.tk_renderWidget.GetRenderWindow().Render()
        self.tk_root.update_idletasks()
        self.tk_root.after(100, self.redraw)

    def shutdown(self):
        self.tk_root.withdraw()
        self.tk_root.destroy()