import pickle
import vtk
import os
import numpy as np
import matplotlib.pyplot as plt
from stack3d.study.components import StudyData, MouseData, TimepointData
import stack3d.formats.unet

def get_time(path):
    
    no_space = path.replace(" ", "")
    day_index = no_space.find("Day")
    day_string = no_space[day_index+3:day_index+5]
    day_string = day_string.replace("_", "")
    return int(day_string)

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
        
        rad_fig, rad_ax = plt.subplots()
        line_fig, line_ax = plt.subplots()
        den_fig, den_ax = plt.subplots()
        
        for eachMouse in eachStudy.mice:
            print "Rendering Mouse: ", eachMouse.raw_data_path
            mouse_dir = study_dir + "/" + eachMouse.raw_data_path
            
            days = []
            radii = []
            lengths = []
            densities = []
            for eachTimePoint in eachMouse.timepoints:
                days.append(get_time(eachTimePoint.raw_data_path))
            if len(days) > 0:
                days, eachMouse.timepoints = (list(t) for t in zip(*sorted(zip(days, eachMouse.timepoints))))            
            
            for kdx, eachTimePoint in enumerate(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()
                    
                    for idx in range(polydata.GetNumberOfPoints()):
                        loc = list(polydata.GetPoint(idx))
                        loc[2] = 0
                        polydata.GetPoints().SetPoint(idx, loc)
                        
                    lines = polydata.GetLines()
                    points = polydata.GetPoints()
                    length = 0.0
                    for jdx in range(polydata.GetNumberOfCells()):
                        pt_ids = vtk.vtkIdList()
                        lines.GetCell(jdx, pt_ids)
                        if pt_ids.GetNumberOfIds() == 2:
                            pt0 = np.array(points.GetPoint(pt_ids.GetId(0)))
                            pt1 = np.array(points.GetPoint(pt_ids.GetId(0)))
                            length += np.linalg.norm(pt0-pt1)
                    lengths.append(length)

                    delaunay = vtk.vtkDelaunay2D()
                    delaunay.SetInputData(polydata)
                    delaunay.Update()

                    com = vtk.vtkCenterOfMass()
                    com.SetInputData(delaunay.GetOutput())
                    com.SetUseScalarsAsWeights(False)
                    com.Update()
                    centre = com.GetCenter()

                    feature = vtk.vtkFeatureEdges()
                    feature.SetBoundaryEdges(1)
                    feature.SetFeatureEdges(0)
                    feature.SetNonManifoldEdges(0)
                    feature.SetManifoldEdges(0)
                    feature.SetInputData(delaunay.GetOutput())
                    feature.Update()
                    
                    av_distance = 0.0
                    for idx in range(feature.GetOutput().GetNumberOfPoints()):
                        loc = feature.GetOutput().GetPoint(idx)
                        av_distance += np.linalg.norm(np.array(centre)-np.array(loc))
                    av_distance /= float(feature.GetOutput().GetNumberOfPoints())
                    densities.append(length/(np.pi*av_distance*av_distance))
                    
                    print "d: ", days[kdx], "c: ", centre, "r: ", av_distance
                    radii.append(av_distance)
                else:
                    radii.append(0)
            rad_ax.plot(days, radii, lw=2.0, label=eachMouse.raw_data_path)
            line_ax.plot(days, lengths, lw=2.0, label=eachMouse.raw_data_path)
            den_ax.plot(days, densities, lw=2.0, label=eachMouse.raw_data_path)

            handles, labels = rad_ax.get_legend_handles_labels()
            rad_ax.legend(handles, labels)
            handles, labels = line_ax.get_legend_handles_labels()
            line_ax.legend(handles, labels)
            handles, labels = den_ax.get_legend_handles_labels()
            den_ax.legend(handles, labels)
            
        rad_fig.savefig(study_dir+"/timeline_radius.png", bbox_inches='tight', dpi=160)
        line_fig.savefig(study_dir+"/timeline_length.png", bbox_inches='tight', dpi=160)
        den_fig.savefig(study_dir+"/timeline_density.png", bbox_inches='tight', dpi=160)

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)