import vtk
from vmtk import pypes

class Surface2d():
    
    def __init__(self, input_path):
        self.input_path = input_path
        self.surface = None
        self.boundaries = None
    
    def generate(self, target_reduction = 0.9994, num_subdivisions = 3, resize_factors = None, clipping_factor = 0.0):
        
        # Convert to VTI
        myArguments = 'vmtkimagereader -ifile ' + self.input_path
        my_pype = pypes.PypeRun(myArguments)
        
        vtk_image = my_pype.GetScriptObject('vmtkimagereader','0').Image
        
        # Size the image if needed
        if resize_factors is not None:
            reduction_filter = vtk.vtkImageResize()
            reduction_filter.SetInput(vtk_image)
            reduction_filter.SetMagnificationFactors(resize_factors[0], resize_factors[1], resize_factors[2])
            reduction_filter.SetResizeMethodToMagnificationFactors()
            reduction_filter.Update()
        
        # Threshold
        threshold = vtk.vtkThreshold()
        if resize_factors is not None:
            threshold.SetInputConnection(reduction_filter.GetOutputPort())
        else:
            threshold.SetInput(vtk_image)
        threshold.ThresholdByUpper(1.0)
        threshold.Update()
        
        # Convert to polydata
        surface = vtk.vtkGeometryFilter()
        surface.SetInputConnection(threshold.GetOutputPort())
        surface.Update()
        
        # Triangulate
        triangle = vtk.vtkTriangleFilter()
        triangle.SetInputConnection(surface.GetOutputPort())
        triangle.Update()
    
        # Decimate
        decimate = vtk.vtkDecimatePro()
        decimate.SetInputConnection(triangle.GetOutputPort())
        decimate.SetTargetReduction(target_reduction)
        decimate.SetFeatureAngle(15.0)
        decimate.Update()
        
        # Do loop subdivision
        su = vtk.vtkLinearSubdivisionFilter()
        su.SetInputConnection(decimate.GetOutputPort())
        su.SetNumberOfSubdivisions(num_subdivisions)
        su.Update()
        
        # Clip the boundaries, recommended to ensure straight inlets and outlets
        if clipping_factor > 0.0:
            bounds = su.GetOutput().GetBounds()
            
            width = bounds[1] - bounds[0]
            height = bounds[3] - bounds[2]
            print bounds[2], bounds[3], height
            
            p1 = vtk.vtkPlane()
            p1.SetOrigin(bounds[0] + clipping_factor*width, 0.0, 0)
            p1.SetNormal(1, 0, 0)
            
            p2 = vtk.vtkPlane()
            p2.SetOrigin(0.0, bounds[2] + clipping_factor*height, 0)
            p2.SetNormal(0, 1, 0)
            
            p3 = vtk.vtkPlane()
            p3.SetOrigin(bounds[1] - clipping_factor*width, 0.0, 0)
            p3.SetNormal(-1, 0, 0)
        
            p4 = vtk.vtkPlane()
            p4.SetOrigin(0.0, bounds[3] - clipping_factor*height, 0)
            p4.SetNormal(0, -1, 0)        
        
            c1 = vtk.vtkClipPolyData()
            c1.SetInputConnection(su.GetOutputPort())
            c1.SetClipFunction(p1)
            c1.SetValue(0.0)
            
            c2 = vtk.vtkClipPolyData()
            c2.SetInputConnection(c1.GetOutputPort())
            c2.SetClipFunction(p2)
            c2.SetValue(0.0)
# 
            c3 = vtk.vtkClipPolyData()
            c3.SetInputConnection(c2.GetOutputPort())
            c3.SetClipFunction(p3)
            c3.SetValue(0.0)
             
            c4 = vtk.vtkClipPolyData()
            c4.SetInputConnection(c3.GetOutputPort())
            c4.SetClipFunction(p4)
            c4.SetValue(0.0)
            
            su = vtk.vtkGeometryFilter()
            su.SetInputConnection(c4.GetOutputPort())
            su.Update()
        
        feature_edges = vtk.vtkFeatureEdges()
        feature_edges.SetInputConnection(su.GetOutputPort())  
        feature_edges.Update()   
        
        clean = vtk.vtkCleanPolyData()
        clean.SetInputConnection(feature_edges.GetOutputPort())
        clean.Update()
         
        triangle2 = vtk.vtkTriangleFilter()
        triangle2.SetInputConnection(clean.GetOutputPort())
        triangle2.Update()
        
        self.surface = su.GetOutput()
        self.boundaries = triangle2.GetOutput()
        return su.GetOutput(), triangle2.GetOutput()
    
    def write(self, output_directory):
        
        surface_writer = vtk.vtkXMLPolyDataWriter()
        surface_writer.SetFileName(output_directory + "/surface.vtp")
        surface_writer.SetInput(self.surface)
        surface_writer.Write() 
        
        surface_writer.SetFileName(output_directory + "/boundaries.vtp")
        surface_writer.SetInput(self.boundaries)
        surface_writer.Write() 
        
class Surface3d():
    
    def __init__(self):
        pass
    
    def generate(self, input_path, output_directory, part_name):
        
        input_path = work_dir + "/vtk_image/itk_tile_0_2_6.vti"
        
        reader = vtk.vtkXMLImageDataReader()
        reader.SetFileName(input_path)
        reader.Update()
         
        pad_filter = vtk.vtkImageConstantPad()
        pad_filter.SetInputConnection(reader.GetOutputPort())
        pad_filter.SetConstant(0)
        pad_filter.SetOutputWholeExtent(reader.GetOutput().GetExtent()[0]-5, 
                                        reader.GetOutput().GetExtent()[1]+5, 
                                        reader.GetOutput().GetExtent()[2]-5, 
                                        reader.GetOutput().GetExtent()[3]+5, 
                                        reader.GetOutput().GetExtent()[4]-5, 
                                        reader.GetOutput().GetExtent()[5]+5)
        pad_filter.Update()
         
        writer = vtk.vtkXMLImageDataWriter()
        writer.SetFileName(work_dir + "/surface/itk_tile_0_2_6_pad.vti")
        writer.SetInputConnection(pad_filter.GetOutputPort())
        writer.Update()    
        
        
        myArguments = 'vmtklevelsetsegmentation -ifile ' + work_dir + "/surface/itk_tile_0_2_6_pad.vti" + ' -ofile ' + output_path+"/itk_tile_0_2_6_ls.vti"
     
        myPype = pypes.PypeRun(myArguments)
         
        myArguments = 'vmtkmarchingcubes -ifile ' + output_path+"/itk_tile_0_2_6_ls.vti" + ' -ofile ' + output_path+"/itk_tile_0_2_6_model.vtp"
     
        myPype = pypes.PypeRun(myArguments)