import os
import inspect
import pickle
from subprocess import call
from PIL import Image
import vtk
import scipy.sparse as sps
from scipy.io import savemat
from stack3d.study.components import StudyData, MouseData, TimepointData
import stack3d.formats.unet


def vtk_to_mat(polydata, path):

    num_points = polydata.GetNumberOfPoints()
    sps_acc = sps.coo_matrix((num_points, num_points))
    data = []
    row_indices = []
    column_indices = []
    lines = polydata.GetLines()
    vtk_data = polydata.GetCellData().GetScalars()
    
    print "np: ", num_points
    for idx in range(lines.GetNumberOfCells()):
        points = vtk.vtkIdList()
        lines.GetCell(idx, points)
        if points.GetNumberOfIds() == 2:
            row_indices.append(points.GetId(0))
            column_indices.append(points.GetId(1))
            data.append(vtk_data.GetTuple1(idx))
    sps_acc = sps_acc + sps.coo_matrix((data, (row_indices, column_indices)), shape=(num_points, num_points))
    savemat(path, {'skeleton': sps_acc})


def skeleton_to_mat(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_distance.mat"))
                    polydata = stack3d.formats.unet.skeleton_to_vtp(time_dir + base_path)
                    vtk_to_mat(polydata, output_path)

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