# Turn a (slightly modified) Slicer color table file into a 
# template model hierarchy scene file for use with the Model Maker.

# To run on an input file (the append call needs a directory where this file is, not including name of this file):
# ./Slicer --python-code "import sys; sys.path.append('/directory/with/MakeColorTableHierarchy.py/'); from MakeColorTableHierarchy import *; ParseColorTableFile('/path/to/hierarchyColors.txt'); exit()"
# which will produce an output of /path/to/hierarchyColors.mrml

import slicer, os, string

# 
# Utility method to add a model hierarchy template and display node
#
def AddColorTableHierarchy(templateScene, name, r, g, b, a, parentID):
  hierarchyNode = templateScene.AddNode(slicer.vtkMRMLModelHierarchyNode())
  hierarchyNode.SetName(name)
  displayNode = templateScene.AddNode(slicer.vtkMRMLModelDisplayNode())
  displayNode.SetColor(r,g,b)
  displayNode.SetOpacity(a)
  hierarchyNode.SetAndObserveDisplayNodeID(displayNode.GetID())
  hierarchyNode.SetParentNodeID(parentID)
  return hierarchyNode.GetID()
  
#
# This method takes in a slightly modified Slicer color table file and 
# produces a template model hierarchy scene file for use with the Model Maker.
# For any colors in the input color table file that are named for structures 
# that are logical children of other structures, add the name of the parent 
# structure as the last element on the child's line (no spaces in the name). 
# This usually requires adding named colors for parent structures, suck as for 
# the skeletal structure. The input file can be out of numerical order, the 
# parent/child relatinships are redone at the end to deal with any orphans 
# encountered during parsing (children who's parent nodes are defined later in the file). 
# For example, the skeleton doesn't have a parent, but the skull has the 
# skeleton for a parent:
# 100 skeleton 180 180 180 255
# 101 skull 180 180 180 255 skeleton
# All structures that don't have parents defined will be children of a new 
# model heirarchy node named for the input file.
# Leading white space is ignored if you wish to indent the lines for ease of reading.
# The output mrml file can be passed to the Model Maker as the Color Hierarchy 
# parameter. It is written in the same directory as the input color table file, with the
# extension .mrml
def ParseColorTableFile(filename):
  tableName = os.path.splitext(os.path.basename(filename))[0]
  print 'Using table name', tableName
  # create a new scene
  templateScene = slicer.vtkMRMLScene()
  # export to a mrml file with the same base name as the input file
  templateFileName = os.path.splitext(filename)[0]+'.mrml'
  templateScene.SetURL(templateFileName)

  # make a top level hierarchy node
  topLevelHierarchyID = AddColorTableHierarchy(templateScene,tableName, 0, 0, 0, 1, None)

  # iterate over lines in the file
  file = open(filename,'r')
  # save parent child relationships
  childParents = {}
  for line in file.readlines():
    words = string.split(line)
    nodeName = words[1]
    r = float(words[2])/255.0
    g = float(words[3])/255.0
    b = float(words[4])/255.0
    a = float(words[5])/255.0
    # is there a parent?
    parentNodeID = topLevelHierarchyID
    if len(words) == 7:
      # is the parent in the scene already?
      parentName = words[6]
      parentNode = templateScene.GetFirstNodeByName(parentName)
      if parentNode != None:
        parentNodeID = parentNode.GetID()
      else:
        # save the name for looking up later
        print 'No parent node named',parentName,'yet, saving for later for child',nodeName
        childParents.setdefault(nodeName,parentName)
    # add a hierarchy
    colorHierarchyID = AddColorTableHierarchy(templateScene,nodeName, r, g, b, a, parentNodeID)
    # print 'Added color hierarchy for node',nodeName,'with id',colorHierarchyID

  # after reading through the file, are there any unparented children?
  numOrphans = len(childParents)
  if numOrphans > 0:
    print 'Fixing',numOrphans,'orphaned hierarchy nodes:\n\t',childParents
    for childName, parentName in childParents.items():
      childNode = templateScene.GetFirstNodeByName(childName)
      parentNode =  templateScene.GetFirstNodeByName(parentName)
      if childNode != None and parentNode != None:
        print 'Setting parent of',childName,'to',parentName
        childNode.SetParentNodeID(parentNode.GetID())


  print 'Saving template to file',templateScene.GetURL()
  templateScene.Commit()