"""

This Source object allows one to create a vtkData object
vtkStructuredGrid, vtkStructuredPoints, vtkRectilinearGrid,
vtkPolyData, vtkUnstructuredGrid or a corresponding reader object and
use MayaVi to visualize it.  This makes it easy to create data on the
fly and then use MayaVi for visualization.  The class handles
save_visualization requests intelligently by dumping the used data to
a file that can be used subsequently.

This code is distributed under the conditions of the GPL (GNU General
Public License).  Look at the file LICENSE included with the
distribution for details.

Copyright (C) 2001, Prabhu Ramachandran.
"""

__author__ = "Prabhu Ramachandran <prabhu_r@users.sf.net>"
__version__ = "$Revision: 1.1 $"
__date__ = "$Date: 2001/11/08 18:45:32 $"

import Base.Objects, Common
import vtkpython, string
import vtkPipeline.vtkMethodParser
import vtkPipeline.ConfigVtkObj

debug = Common.debug

class VtkData (Base.Objects.DataSource):

    """ This Source object allows one to create a vtkData object
    vtkStructuredGrid, vtkStructuredPoints, vtkRectilinearGrid,
    vtkPolyData, vtkUnstructuredGrid or a corresponding reader object
    and use MayaVi to visualize it.  This makes it easy to create data
    on the fly and then use MayaVi for visualization.  The class
    handles save_visualization requests intelligently by dumping the
    used data to a file that can be used subsequently.  """

    def __init__ (self, renwin=None): 
        debug ("In VtkData::__init__ ()")
        Base.Objects.DataSource.__init__ (self)
        self.renwin = renwin
        self.file_name = "No file"
        self.scalar_data_name = "Scalar"
        self.vector_data_name = "Vector"
        self.tensor_data_name = "Tensor"
        self.data_types = {'vtkStructuredGrid': 'STRUCTURED_GRID',
                           'vtkStructuredPoints': 'STRUCTURED_POINTS',
                           'vtkRectilinearGrid': 'RECTILINEAR_GRID',
                           'vtkPolyData': 'POLYDATA',
                           'vtkUnstructuredGrid': 'UNSTRUCTURED_GRID'}
        self.current_type = ''

    def _get_type (self, obj):
        debug ("In VtkData::_get_type ()")
        self.grid_type = ''
        for i in self.data_types.keys ():
            if obj.IsA (i):
                self.grid_type = self.data_types[i]
                self.current_type = i
                break
        if not self.grid_type:
            raise Base.Objects.ParseException, \
                  "Unknown data type: Known data types are one of: %s."\
                  %(self.data_types.keys ())
    
    def initialize (self, obj): 
        """Initialize the object given a valid object."""
        debug ("In VtkData::initialize ()")
        Common.state.busy ()
        if obj.IsA ('vtkDataReader'):
            self.data = obj.GetOutput ()
        else:
            self.data = obj            
        self._get_type (self.data)
        self.update_references ()
        Common.state.idle ()

    def GetOutput (self): 
	"""Get the Data reader's output. """
        debug ("In VtkData::GetOutput ()")
        return self.data

    def save_config (self, file): 
        """Save the configuration to file."""
        debug ("In VtkData::save_config ()")
        # write the data to a vtk file.
        filename = file.name + '.vtk'
        writer_name = self.current_type + 'Writer'
        writer = getattr (vtkpython, writer_name) ()
        writer.SetFileName (filename)
        writer.SetInput (self.data)
        writer.Write ()        
        # now save the necessary data to the mayavi file.
        file.write ("%s, %s\n"%(self.current_type, filename))
        file.write ("%s, %s, %s\n"%(self.scalar_data_name,
                                    self.vector_data_name,
                                    self.tensor_data_name))
        p = vtkPipeline.vtkMethodParser.VtkPickler ()
        p.dump (self.data, file)

    def load_config (self, file): 
        """Load the saved objects configuration from a file."""
        debug ("In VtkData::load_config ()")
        val = string.replace (file.readline ()[:-1], ' ', '')
        val = string.split (val, ',')
        typ, filename = val
        val = string.replace (file.readline ()[:-1], ' ', '')
        val = string.split (val, ',')
        s_n, v_n, t_n = val
        self.scalar_data_name = s_n
        self.vector_data_name, self.tensor_data_name = v_n, t_n
        # set type info.
        self.current_type = typ
        self.grid_type = self.data_types[typ]
        # read data from vtk file.
        reader_name = self.current_type + 'Reader'
        reader = getattr (vtkpython, reader_name) ()
        reader.SetFileName (filename)
        reader.Update ()        
        self.data = reader.GetOutput ()
        # now read the config for the object
        p = vtkPipeline.vtkMethodParser.VtkPickler ()
        p.load (self.data, file)
        self.update_references ()
        
    def configure (self, master=None):
        debug ("In VtkData::configure ()")
        c = vtkPipeline.ConfigVtkObj.ConfigVtkObj (self.renwin)
        c.configure (master, self.data)
        c.root.transient (master)

    def set_scalar_data_name (self, name):
	"""Set the name of the scalar data attribute being used."""
        debug ("In VtkData::set_scalar_data_name ()")
	self.scalar_data_name = name
        self.update_references ()

    def set_vector_data_name (self, name):
	"""Set the name of the vector data attribute being used."""
        debug ("In VtkData::set_vector_data_name ()")
	self.vector_data_name = name
        self.update_references ()

    def set_tensor_data_name (self, name):
	"""Set the name of the tensor data attribute being used."""
        debug ("In VtkData::set_tensor_data_name ()")
	self.tensor_data_name = name
        self.update_references ()
