#!/home/brian/VEnvs/graphPaper/bin/python
from vasctrees.SkeletonGraph import SkeletonGraph
import vasctrees.SkeletonGraph as skgs
import sys
import scipy.ndimage as ndi
import numpy as np
import SimpleITK as sitk

def readImage(fname):
    img = sitk.ReadImage(fname)
    aimg = sitk.GetArrayFromImage(img)
    descriptors = {}
    descriptors["scale"] = img.GetSpacing()
    descriptors["origin"] = img.GetOrigin()
    descriptors["orientation"] = img.GetDirection()
    return aimg, descriptors

simg,descr = readImage( sys.argv[1] )

try:
    simg
except:
    sys.exit("read of skeleton image failed")

sg = SkeletonGraph(img=simg,
                   spacing=descr['scale'],
                   origin=descr['origin'],
                   orientation=descr['orientation'],
		   label = sys.argv[1])

##print descr
#print sg.spacing, sg.origin, sg.orientation

oimg,tmp = readImage( sys.argv[2])
try:
    oimg
except:
    sys.exit("read of mask image failed")
sg.setOriginalImage(oimg)
sg.getGraphsFromSkeleton(verbose=False)
sg.setLargestGraphToCurrentGraph()
sg.graphs.keys()
sg.findEndpointsBifurcations()
endp = [n for n in sg.cg.nodes() if sg.cg.degree(n)==1]
endpa = np.array(endp)
medianx = np.median(endpa[:,0])
endp.sort(key=lambda n: abs(n[0]-medianx))
root = endp[0]
#sg.viewGraph()
sg.setRoot(root,key="og_medianx")
sg.traceEndpoints(key='og_medianx')
ogkey = sg.getLargestOrderedGraphKey()
sg.deleteDegree2Nodes(ogkey)
sg.prunePaths(ogkey)
sg.deleteDegree2Nodes(ogkey)
sg.fitEdges(key=ogkey)
#print "Define orthogonal planes"
sg.defineOrthogonalPlanes(ogkey)
# Now get surface points of original image to map to the centerlines
dfe = ndi.distance_transform_cdt(oimg)
points_toMap = np.array(np.nonzero(np.where(dfe==1,1,0))[::-1]).transpose().astype(np.int32)
#print "mapVoxelsToGraph"
sg.mapVoxelsToGraph(points_toMap,ogkey)
#print "assignMappedPointsToPlanes"
sg.assignMappedPointsToPlanes(ogkey)
sg.saveCompressedGraphs(sys.argv[3])
