from vmtk import vmtkscripts
from vtk import *
import os

def vmtksurfacereader(filename):
    reader = vmtkscripts.vmtkSurfaceReader()
    reader.InputFileName = filename
    reader.Execute()
    return reader.Surface

def vmtksurfacewriter(polydata, filename):
    writer = vmtkscripts.vmtkSurfaceWriter()
    writer.Surface = polydata
    writer.OutputFileName = filename
    writer.Execute()



os.chdir('D:/temp')

input = vmtksurfacereader('patchedbranch.vtp')

numberofpoints = input.GetNumberOfPoints()
points = vtk.vtkPoints()
points.SetNumberOfPoints(numberofpoints)    

for i in range(numberofpoints):
    pointx = input.GetPointData().GetArray('AngularMetric').GetValue(i)
    pointy = input.GetPointData().GetArray('StretchedMapping').GetValue(i)
    pointz = 0.0
    point = (pointx, pointy, pointz)
    points.SetPoint(i,point)

output = vtk.vtkPolyData()
output.DeepCopy(input)
output.SetPoints(points)

vmtksurfacewriter(output, 'patchedbranchflat.vtp')    

