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


def skeleton_to_vtp(work_dir, study_collection):

    for eachStudy in study_collection:
        print "Converting Study: ", eachStudy.raw_data_path
        study_dir = work_dir + "/" + eachStudy.raw_data_path
        for eachMouse in eachStudy.mice:
            print "Converting Mouse: ", eachMouse.raw_data_path
            mouse_dir = study_dir + "/" + eachMouse.raw_data_path
            for eachTimePoint in eachMouse.timepoints:
                print "Converting Time: ", eachTimePoint.raw_data_path
                time_dir = mouse_dir + "/" + eachTimePoint.raw_data_path
                base_path = "/analysis_analysis/skeleton.pkl"
                if os.path.isfile(time_dir + base_path):
                    output_path = os.path.dirname(os.path.realpath(time_dir + "/analysis_analysis/skeleton.vtp"))
                    polydata = stack3d.formats.unet.skeleton_to_vtp(time_dir + base_path)
                    writer = vtk.vtkXMLPolyDataWriter()
                    writer.SetInputData(polydata)
                    writer.SetFileName(output_path + "/skeleton.vtp")
                    writer.Write()

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