"""Utilities for creating VTK datasets from Numeric arrays."""
import Numeric as Nm
import vtk

# Map numeric typecode to VTK array class
nm_code_to_vtk = {'d':vtk.vtkDoubleArray,
                  'f':vtk.vtkFloatArray,
                  'l':vtk.vtkIntArray,
                  '1':vtk.vtkCharArray,
                  's':vtk.vtkShortArray,
                  'i':vtk.vtkIntArray}


def nm_to_vtk_array(nm_array,nm_type=None,fortranorder=True):
    """Given a Numeric array, create a VTK array from it. The type of
    the resulting vtk array will correspond to nm_type (e.g.
    Nm.Float32 ==> vtkFloatArray). If nm_type is not provided, the
    type of nm_array is used.
    """
    if nm_type is not None:
        outdata=nm_array.astype(nm_type)
    else:
        outdata=nm_array
        nm_type=nm_array.typecode()
        
    if fortranorder:
        raw=Nm.transpose(nm_array).tostring()
    else:
        raw=nm_array.tostring()
    # Generate the appropriate type of VTKArray from the numeric typecode.
    temp_array=nm_code_to_vtk[nm_type]()
    
    nelems=Nm.size(nm_array) 
    temp_array.SetNumberOfValues(nelems)
    temp_array.SetVoidArray(raw, nelems,1)

    # One problem is that VTK will not increase the reference count of
    # the raw data, so if we just return temp_array, the storage holding
    # the data will be reclaimed, and we'll get a lovely core dump. Therefore
    # we create a deep copy of the data, for which VTK will allocate (and handle)
    # memory. Slightly inefficient, but it saves from having to keep track of
    # the raw data separately.
    result=nm_code_to_vtk[nm_type]()
    result.DeepCopy(temp_array)
    return result


def nm_to_vtk_structured_points(nm_array,fortranorder=True):
    """ Generate a vtkStructuredPoints dataset from a Numeric array.
    Using vtkImageImporter (new in vtk5.x?) should be prefered over
    this method.
    Returns a 2-tuple where the first element is a VTK dataset with the
    
    The result will be a dataset with the same dimensions as the input,
    with default origin and spacing
    """
    result=vtk.vtkStructuredPoints()
    shp=nm_array.shape
    if len(shp)==2:
	shp=(shp[0],shp[1],1)
    if (not fortranorder):
        shp=shp[::-1]
    
    result.SetDimensions(shp)
    array=nm_to_vtk_array(nm_array,fortranorder=fortranorder)
    result.GetPointData().SetScalars(array)
    return result

def nm_to_vtk_rectilinear_grid(xc,yc,zc,nm_array,fortranorder=True):
    """ Generate a vtkRectilinearGrid dataset from a bunch of Numeric
    arrays. xc,yc, and zc  contain 1D arrays of x, y, and z
    coordinate values, whereas nm_array contains a 3D array of values.
    The length of each coordinate array should equal the corresponding
    extent of nm_array (i.e. (xc.size(),yc.size(),zc.size()) == nm_array.shape())
    This method also accepts 2D arrays as input, in which case zc is ignored.
    """
    result=vtk.vtkRectilinearGrid()
    shp=nm_array.shape
    if len(shp)==2:
        result.SetDimensions(shp[0],shp[1],1)
    else:
   	result.SetDimensions(shp)
    xdata=nm_to_vtk_array(xc)
    result.SetXCoordinates(xdata)
    
    ydata=nm_to_vtk_array(yc)
    result.SetYCoordinates(ydata)
    
    if len(shp)==3:
	zdata=nm_to_vtk_array(zc)
        result.SetZCoordinates(zdata)
	zraw=0
    
    array=nm_to_vtk_array(nm_array,fortranorder=fortranorder)
    result.GetPointData().SetScalars(array)
    return result

def add_textures(ds,u,v=None,fortranorder=True):
    """ Given a vtkStructuredPoints or vtkRectilinearGrid dataset,
    and one or two numeric arrays of the same extents, AddTextures
    adds the arrays as the u and v texture coordinates of the data
    (useful for using a texture coordinate as a lookup into a colormap).
    """
    dims=ds.GetDimensions()
    sz=reduce(lambda x,y:x*y,dims,1)
    if ((u and u.shape!=dims) or (v and v.shape!=dims)):
        print >>sys.stderr,"Texture data has improper extents"
        return
    # Always create 2D texture coordinates, even if only one is supplied,
    # because some tools like Maya don't seem to handle 1D texture coordinates
    # properly.
    arr=Nm.zeros(2*sz,Nm.Float32)
    if fortranorder:
        if u:
            arr[::2]=Nm.ravel(Nm.transpose(u.astype(Nm.Float32)))
        if v:
            arr[1::2]=Nm.ravel(Nm.transpose(v.astype(Nm.Float32)))
    else:
        if u:
            arr[::2]=Nm.ravel(u.astype(Nm.Float32))
        if v:
            arr[1::2]=Nm.ravel(v.astype(Nm.Float32))
    raw=arr.tostring()
    tc=vtk.vtkFloatArray()
    tc.SetNumberOfComponents(2)
    tc.SetNumberOfTuples(sz)
    tc.SetVoidArray(raw ,2*sz,1)
    tc_permanent=vtk.vtkFloatArray()
    tc_permanent.DeepCopy(tc)
    ds.GetPointData().SetTCoords(tc_permanent)
    
def add_field_data(ds,scalar_data,scalar_name,scalar_type=None,fortranorder=True):
    """ AddFieldData adds a field to an existing vtk dataset (StructuredPoints or
    RectilinearGrid). scalar_data is a Numeric array, and must have the same shape
    as the dataset. By default, 32 bit floats are used.
    """
    dims=ds.GetDimensions()
    sz=reduce(lambda x,y:x*y,dims,1)
    if (scalar_data.shape!=dims):
        print >>sys.stderr,"Scalar data has improper extents"
        return

    sarr=nm_to_vtk_array(scalar_data,scalar_type,fortranoder=fortranorder)
    sarr.SetName(scalar_name)
    ds.GetPointData().AddArray(sarr)