#!/usr/bin/python
#dv3d_vtkTools_generateVTKSurface.py
#uses vtk contourin to contruct a surface from loaded data
#DV3D refactor 2
#AG 25/07/2013

import vtk
import colorsys


from dv3d_data_AddSurfaceToWindow import *


def dv3d_vtkTools_generateVTKSurfaceFromData(the_parent_window, dataIndex, extractionValue, surfacesToGetIndex, smoothingIterations, samplingFrequency, decimationFactor, tgtColor=(1,0,0), tgtOpacity=1.0, thisItemIndexInGroup=0, groupCount=0):
    #extract surface type information from volumeteric data usinf sampling and contour fitering
    # then call a generic script that takes the ouptut and converts it to a vtk surface
    #  then add details to the dv3d structure
    
    tpw = the_parent_window

    # get the volume
    vol = tpw.Planes[dataIndex].volumeData
    
    theVTKVolume = vol.GetOutput() #get the underlying volume
    theVTKVolume.Update() #actually load the data

    # sampling code 
    # - a value of 1 = sample every adjacent pixel/voxel
    # - a value of 3 = sample every 3rd pixel/voxel - effectively makes the matrix more course
    shrink = vtk.vtkImageShrink3D()
    shrink.SetShrinkFactors(samplingFrequency, samplingFrequency, samplingFrequency)
    shrink.SetInput(theVTKVolume)
    shrink.AveragingOn()

    # actual surface extraction     
    isoSurfaceExtractor = vtk.vtkContourFilter() #used to generate a suface by connecting scalars of equal value
    isoSurfaceExtractor.UseScalarTreeOn()
    isoSurfaceExtractor.SetInput(shrink.GetOutput())    
    isoSurfaceExtractor.SetValue(0, extractionValue) #sets the 'range' of the dots to connect
    isoSurfaceExtractor.ComputeGradientsOn()
    isoSurfaceExtractor.UseScalarTreeOn()
    
    thePolyData = isoSurfaceExtractor.GetOutput()
    thePolyData.Update()


    smooth = vtk.vtkSmoothPolyDataFilter()
    smooth.SetNumberOfIterations(smoothingIterations)
    smooth.SetRelaxationFactor(decimationFactor)
    smooth.FeatureEdgeSmoothingOff()
    smooth.SetInput(thePolyData)
    smooth.Update()



    #now process and convert the vertex/line information into something we can see
    theSurface = generateTheVTKSurface(tpw, smooth.GetOutput(), surfacesToGetIndex, smoothingIterations, decimationFactor, tgtColor, tgtOpacity)

    #let's keep track of where the surface came from
    setattr(theSurface, 'dataLabel', 'from %s at %0.4f' %(tpw.Planes[dataIndex].dataLabel, extractionValue))
    
    #let's keep track of whether its a single surface or part of a group
    if thisItemIndexInGroup == 0 and groupCount <= 1: # is a single surface
        setattr(theSurface, 'memberOfGroup',None)
        setattr(theSurface, 'groupItemNumber',None)
    elif groupCount > 1: #at least two surfaces so must be a group
        setattr(theSurface, 'memberOfGroup', tpw.uniqueSurfaceGroupCount) #NB NOT zero indexed here
        setattr(theSurface, 'groupItemNumber',thisItemIndexInGroup) # order in current group

    #now pass it to the routine that adds it with the DV3D structure
    dv3d_data_AddSurfaceToWindow(tpw, theSurface, thisItemIndexInGroup, groupCount)



def dv3d_vtkTools_generateFromLoadedFile(the_parent_window, fileName):

    tpw = the_parent_window #for brevity

    #load the data 
    #TODO - add different file reader types
    
    reader = vtk.vtkPolyDataReader()
    reader.SetFileName(fileName)
    
    ##TODO - re-insert this: freesurfer support     
    ##correct freesurfer offset if required
    #xfm = vtk.vtkTransform()
    ##if frame_parent.applyVTKxfm == 1:
    #if vertex_file[-6:-5] != 'o':    
    #    xfm.Translate(0, 0, 0)
    #    #xfm.Translate(88, 128, 128)
    #    #xfm.Scale(-1.0, 1.0, 1.0)
    #else:
    #    xfm.Translate(0, 0, 0)
    #    xfm.Scale(1.0, 1.0, 1.0)
    #    frame_parent.applyVTKxfm = 1
    #
    #transformer1 = vtk.vtkTransformPolyDataFilter()
    #transformer1.SetInput(try_me.GetOutput())
    #transformer1.SetTransform(xfm)


    thePolyData = reader.GetOutput()
    thePolyData.Update()

    theSurface = generateTheVTKSurface(tpw, thePolyData, 0, 0, 0) #load as is for now
    #TODO - allow passing or args whne loading a surface?

    #let's keep track of where the surface came from
    setattr(theSurface, 'dataLabel', fileName)

    #now pass it to the routine that adds it with the DV3D structure
    dv3d_data_AddSurfaceToWindow(tpw, theSurface)


def generateGroupSurfaces(the_parent_window, volumeDataIndex, extractVals, extractOpacities, extractHSVs, surfacesToGetIndex, smoothingIterations, decimationFactor, samplingFrequency):
    #this function generates a vtk surface for each item in the list passed to it: called by multi-surface functions
    tpw = the_parent_window #for brevity
    groupCount = len(extractVals)

    for i in range(len(extractVals)):
        thisItemIndexInGroup = i
        tgtColor = colorsys.hsv_to_rgb(extractHSVs[i],1,1)
        dv3d_vtkTools_generateVTKSurfaceFromData(tpw, volumeDataIndex, extractVals[i], surfacesToGetIndex, smoothingIterations, samplingFrequency, decimationFactor, tgtColor, extractOpacities[i], thisItemIndexInGroup, groupCount)



#def generateTheVTKSurface(the_parent_window, theVTKPolyDataObject, surfacesToGetIndex, smoothingIterations, decimationFactor):
def generateTheVTKSurface(the_parent_window, theVTKPolyDataObject, surfacesToGetIndex, smoothingIterations, decimationFactor, surfaceColor=(1,0,0), surfaceOpacity=1.0):

    tpw = the_parent_window #for brevity

    # a generic script that takes the vtkPolydata (generated or laoded) and converts it to a vtk surface    
    smooth = vtk.vtkSmoothPolyDataFilter()
    smooth.SetNumberOfIterations(smoothingIterations)
    smooth.SetRelaxationFactor(decimationFactor)
    smooth.FeatureEdgeSmoothingOff()
    smooth.SetInput(theVTKPolyDataObject)
    smooth.Update()

    print 'orig'
    print theVTKPolyDataObject.GetNumberOfPoints()

    print 'new'
    print smooth.GetOutput().GetNumberOfPoints()


    surfNormals = vtk.vtkPolyDataNormals() #required for correct display 
    ##surfNormals.SetInput(pOut)
    surfNormals.SetInput(smooth.GetOutput())
    surfNormals.SetInput(theVTKPolyDataObject)
    surfNormals.SetFeatureAngle(60.0)

    conn = vtk.vtkPolyDataConnectivityFilter() #actually connect the dots now     
    conn.SetInput(surfNormals.GetOutput())
    if surfacesToGetIndex == 1:
        conn.SetExtractionModeToLargestRegion() #show only the largest connected surface
    else:
        conn.SetExtractionModeToAllRegions()
    conn.Update()
    
    # since the process above generates surfaces with (typically) 600k vertices,
    # this option allows us to downsample the surface structure. This is useful
    # for speeding up rendering (we spend a bitmore time here ~10 seconds)
    # TODO - option to display either the decimated / full surface
    deci = vtk.vtkDecimatePro() 
    deci.SetInput(conn.GetOutput())
    deci.PreserveTopologyOn()
    deci.SetTargetReduction(decimationFactor) #target decimation factor -> 0.9 = 90% reduction in number of vertices
    deci.Update()

    #now pass our surface to the tool that draws(or 'maps') the data to the
    #screen
    surfMapper = vtk.vtkPolyDataMapper()
    surfMapper.ImmediateModeRenderingOn()
    surfMapper.SetInput(deci.GetOutput())
    surfMapper.ScalarVisibilityOff() 
    surfMapper.Update()

    # and add this data to an actor class
    surf = vtk.vtkActor()
    surf.GetProperty().SetLineWidth(1)
    surf.GetProperty().SetColor(surfaceColor[0], surfaceColor[1], surfaceColor[2])
    surf.GetProperty().SetOpacity(surfaceOpacity)
    surf.SetMapper(surfMapper)

    return surf

