import pickle
import vtk
import os
from stack3d.study.components import StudyData, MouseData, TimepointData
import stack3d.formats.unet


def render_skeleton(work_dir, study_collection):

    for eachStudy in study_collection:
        print "Rendering Study: ", eachStudy.raw_data_path
        study_dir = work_dir + "/" + eachStudy.raw_data_path
        for eachMouse in eachStudy.mice:
            print "Rendering Mouse: ", eachMouse.raw_data_path
            mouse_dir = study_dir + "/" + eachMouse.raw_data_path
            for eachTimePoint in eachMouse.timepoints:
                print "Rendering Time: ", eachTimePoint.raw_data_path
                time_dir = mouse_dir + "/" + eachTimePoint.raw_data_path
                base_path = "/analysis_analysis/skeleton.vtp"
                if os.path.isfile(time_dir + base_path):
                    reader = vtk.vtkXMLPolyDataReader()
                    reader.SetFileName(time_dir + base_path)
                    reader.Update()
                    polydata = reader.GetOutput()
                    
                    cell_to_point = vtk.vtkCellDataToPointData()
                    cell_to_point.SetInputData(polydata)
                    cell_to_point.Update()

                    tubes = vtk.vtkTubeFilter()
                    tubes.SetInputData(cell_to_point.GetOutput())
                    tubes.SetNumberOfSides(6)
                    tubes.SetRadius(1.0)
                    tubes.SetRadiusFactor(50.0)
                    tubes.SetVaryRadiusToVaryRadiusByScalar()
                    tubes.Update()

                    mapper = vtk.vtkPolyDataMapper()
                    mapper.SetInputData(tubes.GetOutput())
                    mapper.ScalarVisibilityOff()

                    actor = vtk.vtkActor()
                    actor.SetMapper(mapper)
                    actor.GetProperty().SetColor(1, 0, 0)
                    
                    textActor = vtk.vtkTextActor()
                    textActor.SetTextScaleModeToProp()
                    textActor.SetDisplayPosition(400, 50)
                    
                    inp_string = "Study: " + eachStudy.raw_data_path
                    inp_string += " | Mouse: " + eachMouse.raw_data_path
                    inp_string += " | Time: " + eachTimePoint.raw_data_path
                    
                    textActor.SetInput(inp_string)
                    tprop = textActor.GetTextProperty()
                    tprop.SetFontSize(24)
                    tprop.SetFontFamilyToArial()
                    tprop.SetJustificationToCentered()
                    tprop.BoldOn()
                    tprop.ItalicOn()
                    tprop.ShadowOn()
                    tprop.SetColor(0, 0, 1)

                    cubesource = vtk.vtkCubeSource()
                    cubesource.SetBounds(0, 4000, 0, 4000, 0, 70)
                    cubesource.Update()
                    cube_mapper = vtk.vtkPolyDataMapper()
                    cube_mapper.SetInputConnection(cubesource.GetOutputPort())
                    cube_mapper.ScalarVisibilityOff()
                    cube_actor = vtk.vtkActor()
                    cube_actor.SetMapper(cube_mapper)
                    cube_actor.GetProperty().SetColor(0, 0, 0)
                    cube_actor.GetProperty().SetOpacity(0.1)

                    # Create a renderer, render window, and interactor
                    renderer = vtk.vtkRenderer()
                    renderer.SetBackground(1, 1, 1)
                    renderer.AddActor2D(textActor)
                    #renderer.GetActiveCamera().Zoom(1.0)
                    renderWindow = vtk.vtkRenderWindow()
                    renderWindow.SetSize(1200, 800)
                    renderWindow.SetAlphaBitPlanes(1)
                    renderWindow.AddRenderer(renderer)
#                     renderWindowInteractor = vtk.vtkRenderWindowInteractor()
#                     renderWindowInteractor.SetRenderWindow(renderWindow)
                    # Add the actors to the scene
                    renderer.AddActor(actor)
                    renderer.AddActor(cube_actor)
                    renderer.ResetCameraClippingRange()
                    # Render and interact
                    renderWindow.Render()
                    #renderWindowInteractor.Start()

                    window_to_image_filter = vtk.vtkWindowToImageFilter()
                    window_to_image_filter.SetInput(renderWindow)
                    window_to_image_filter.SetMagnification(1);
                    window_to_image_filter.SetInputBufferTypeToRGBA()
                    window_to_image_filter.ReadFrontBufferOff()
                    writer = vtk.vtkTIFFWriter()
                    writer.SetInputConnection(window_to_image_filter.GetOutputPort())
                    writer.SetFileName(time_dir + "/analysis_analysis/rendering.tiff")
                    writer.Write()

work_dir = "/scratch/jgrogan/stack-working/study/"
f = open(work_dir + "/study_collection.p", 'r')
study_collection = pickle.load(f)
render_skeleton(work_dir, study_collection)