#!/usr/bin/env python


import vtk
from vtk.util import numpy_support
from importlib import reload
import numpy as np
import os
import mathutils

try:
    import bpy
    inside_blender = True
except:
    inside_blender = False

if (inside_blender):
    import VTKBlender
    reload(VTKBlender)

# Global, persistent data
#save_frames = True
frame_change = None
dataRoot = '/mydata/'
dataFile = 'testData.mha'


def main():
    print('Starting Main Function...')
    if (not inside_blender):
        return
    old_pre_count = len(bpy.app.handlers.frame_change_pre)
    for i in range(old_pre_count):
        print('Removing frame_change_pre handle {}'.format(i))
        bpy.app.handlers.frame_change_pre.pop()
    bpy.app.handlers.frame_change_pre.append(frame_change_pre)

def frame_change_pre(scene):
    print('My frame change pre function')
    global frame_change, npdata
    if (frame_change == None):
        print('initializing new frame changer')
        frame_change = FrameChange(scene, dataFile)
    else:
        print('using old FrameChange Object')
        frame_change.update(scene)

class FrameChange:
    def __init__(self, scene, dataFile):
        self.vtkDataReader = vtk.vtkMetaImageReader()
        self.vtkDataReader.SetFileName( os.path.join(dataRoot, dataFile) )
        self.vtkDataReader.Update()

        self.level = 1
        self.contours = vtk.vtkContourFilter()
        self.stlWriter = vtk.vtkSTLWriter()

        self.init_frame(scene)
        self.init_lut()
        self.init_mapper()
        self.init_mesh(scene)
        print('Initialized Frame Changern' + 15*'-')


    def update(self, scene):
        print('Updating scene')
        global save_frames
        self.init_frame(scene)
        self.update_contours()
        self.mapper.Update()
        #if save_frames == True:
            #self.stlWriter.SetFileName(os.path.join(dataRoot, 'tmp',str(self.level) + '_contour.stl'))
            #print('file to write: {}'.format(self.stlWriter.GetFileName()))
            #self.contours.Update()
            #self.stlWriter.SetInputConnection( self.contours.GetOutputPort() )
            #self.stlWriter.Write()
        VTKBlender.PolyDataMapperToBlender.convert(self.mapper, self.me)
        self.me.update(calc_edges = True)
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.remove_doubles()
        bpy.ops.mesh.normals_make_consistent(inside=False)
        bpy.ops.object.mode_set(mode='OBJECT')

        self.smooth_faces()

    def init_frame(self, scene):
        print('Initializing frame')
        self.level +=1


    def smooth_faces(self):
        print ('Smoothing polygon faces')
        for face in self.me.polygons:
            face.use_smooth = True

    def init_mapper(self):
        print('Initializing mapper')
        self.contours.SetInputData( self.vtkDataReader.GetOutput() )
        self.update_contours()
        #self.mapper = get_basic_mapper(self.contours, self.lut)
        #self.stlWriter.SetFileName(os.path.join(dataRoot,'tmp',str(self.level) + '_contour.stl'))
        #print('file to write: {}'.format(self.stlWriter.GetFileName()))
        #self.contours.Update()
        #self.stlWriter.SetInputConnection( self.contours.GetOutputPort() )
        #self.stlWriter.Write()

    def update_contours(self):
        print('Updating isosurfaces level')
        self.contours.GenerateValues(2, self.level, self.level+0.5 )


    def init_mesh(self,scene):
        print('Initializing mesh')
        self.me = bpy.data.meshes.new('mySurface')
        self.object = bpy.data.objects.new("mySurface",self.me)
        self.object.location = mathutils.Vector((0, 0, 0))
        scene.objects.link(self.object)
        VTKBlender.PolyDataMapperToBlender.convert(self.mapper, self.me)
        self.me.update(calc_edges = True)
        scene.objects.active = self.object
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.remove_doubles()
        bpy.ops.mesh.normals_make_consistent(inside=False)
        bpy.ops.object.mode_set(mode='OBJECT')


    def init_lut(self):
        print('Initializing lut')
        self.lut = vtk.vtkLookupTable()
        self.lut.SetNumberOfColors( 10 )
        self.lut.SetHueRange(0.667,0.0);
        self.lut.Build()


def get_basic_mapper(ob, lut):
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(ob.GetOutputPort())
    mapper.SetLookupTable(lut)
    return mapper

main()