Interactive/AutoGroupCustom.py

##
# (C) Copyright 2012 Autodesk, Inc.  All rights reserved.
##
## Use of this software is subject to the terms of the Autodesk license
## agreement provided at the time of installation or download, or which
## otherwise accompanies this software in either electronic or hard copy
## form.
##

__all__ = ['AutoGroupCustom', 'instantiate']

from Message             import Message
from MessageInterpreter  import MessageInterpreter
from MessageRegistry     import theMessageRegistry
from Modes               import ViewportId
from SceneGraphUtilities import CastToOriginalType, \
                                CollectLeafNodes, \
                                CollectModelLeafNodes, \
                                CollectNodes, \
                                GetModelImport, \
                                FindLowestRealParent, \
                                kStopOnSuccess
from UserCustomization   import UserCustomBase, CustomInfo
from Utilities           import AlphanumericSort
from awSupportUtilities  import FindFile

import ModelIO, wx

AUTO_GROUP_Doc = \
[(
"""
Groups [and collapses] based on Revit category, family, type, and material.
"""
),
[[( 'collapse' ), ( 'Whether to also collapse the new groups' )]
 ]
]

theMessageRegistry.register( 'AUTO_GROUP', ( bool, ), Message.kUndoable, AUTO_GROUP_Doc )
theMessageRegistry.register( 'AUTO_GROUP_COLLAPSE', (), Message.kInternal, None )


# Created for the Pilot, not production quality code.
# - works on the entire model in the scene;
# - automatically groups the flat node hierarchy in a Revit 
#   fbx file by category, family, type, and material;
# - additionally, the user can request the collapse of the created groups.
#
class AutoGroupCustom ( UserCustomBase ):
    def __init__( self ):
        self.__myInterpreter = LocalInterpreter()

    def getInterpreter(self,isInteractive):
        return (self.__myInterpreter if isInteractive else None)

    def appendMenuItems( self, id, menu ):
        if 'View' == id:
            menu.AppendSeparator()
            ag = wx.MenuItem( menu, wx.NewId(),
                              _('Group by Category, Family, Type, Material' ) )
            agi = wx.Image( FindFile( 'image', "Group.png" ) )
            if agi and agi.Ok():
                bitmap = agi.ConvertToBitmap()
                ag.SetBitmap( bitmap )
            self.__myAutoGroupId = menu.appendMenuItem( ag,self.__onAutoGroup )

            ac = wx.MenuItem( menu, wx.NewId(),
                              _('Group and Combine by Category, Family, Type, Material' ) )
            aci = wx.Image( FindFile( 'image', "Combine.png" ) )
            if aci and aci.Ok():
                bitmap = aci.ConvertToBitmap()
                ac.SetBitmap( bitmap )
            self.__myAutoCollapseId = menu.appendMenuItem( ac, self.__onAutoCollapse )

    def enableMenuStates( self, id, enableStates ):
        if 'View' == id:
            enableStates[self.__myAutoGroupId] = self.__myInterpreter.isSceneOpen()
            enableStates[self.__myAutoCollapseId] = self.__myInterpreter.isSceneOpen()

    def __onAutoGroup( self, event ):
        self.sendMessage( 'AUTO_GROUP', ( False, ) )

    def __onAutoCollapse( self, event ):
        self.sendMessage( 'AUTO_GROUP', ( True, ) )

    
class LocalInterpreter( MessageInterpreter ):
    def __init__( self ):
        MessageInterpreter.__init__( self )
        self.__myModels = None
        self.__myActiveStageId = ''
        self.__myTransactionId = None
        self.__myGroupIds = []
        self.__myGroupLabels = {}

    def isSceneOpen( self ):
        return self.__myModels is not None

    @staticmethod
    def __getRevitObjectModelInfo( target ):
        # Start with the target node and then traverse upwards until
        # we find an ancestor node that has the info available.
        # Note, we use the original nodes so that we can get the 
        # original metadata in case this node was reparented.
        #
        while target:
            metaData = target.getMetaData()

            if metaData:
                category = metaData.getObjectCategory()
                if category != '':
                    family = metaData.getObjectFamily()
                    if family != '':
                        type = metaData.getObjectType()
                        if type != '':
                            return ( category, family, type )

            target = target.getParent( 0 )
        
        return None

    def __groupNodes( self, nodes, groupLabel ):
        """
        Creates a group for the given nodes, using the given groupLabel, if the
        group does not already exist.  
        """
        if len( nodes ) > 1:
            # First attempt to find a common ancestor for these nodes.  If one 
            # is found, we'll just relabel it.
            #
            alreadyGrouped = False
            ancestor = FindLowestRealParent( nodes )
            if ancestor is not None and ancestor.get() != self.__myModels.root.get():
                if set( nodes ) == set( CollectLeafNodes( ancestor, (), idsOnly=True ) ):
                    ancestorId = ancestor.getUniqueId()
                    self.__myGroupIds.append( ancestorId )
                    
                    if ancestor.getLabel() != groupLabel:
                        self.sendMessage( 'NODE_RENAME', ( ( ancestorId, ), groupLabel ) )

                    alreadyGrouped = True
            
            # If no ancestor is found, create one, and remember the label for later.
            #
            if not alreadyGrouped:
                nodes = tuple( nodes )
                self.sendMessage( 'GROUP_NODES', ( tuple( nodes ), ) )
                self.__myGroupLabels[nodes] = groupLabel
            
            return True
        return False

    def APPLICATION_CLOSE_SCENE( self, message ):
        self.__myModels = None
        self.__myActiveStageId = ''

    def SET_DOCUMENT( self, message ):
        document, = message.data
        self.__myModels = document.get( ModelIO.id )

    def STAGE_ACTIVE_LIST( self, message ):
        ( viewportIdsToStageIds, ) = message.data

        # Note, we currently only work on nodes from the current stage.  
        #
        self.__myActiveStageId = viewportIdsToStageIds.get( ViewportId.kLeft, '' )

    def NODE_CREATED( self, message ):
        # If this node was not created by AUTO_GROUP, there's nothing to do.
        #
        if message.transactionId != self.__myTransactionId:
            return

        # Remember this group's id as we may be collapsing it later (if requested).
        #
        ( groupId, parentId, insertedIndex ) = message.data
        self.__myGroupIds.append( groupId )
        
        # Search for group's label and relabel it.
        #
        group = self.__myModels[groupId]
        child = group.getChild( 0 )
        childId = child.getUniqueId()

        for ( nodeIds, groupLabel ) in self.__myGroupLabels.iteritems():
            if childId in nodeIds:
                break
        
        del self.__myGroupLabels[nodeIds]
        self.sendMessage( 'NODE_RENAME', ( ( groupId, ), groupLabel ) )

    def AUTO_GROUP_COLLAPSE( self, message ):
        if self.__myTransactionId != message.transactionId or len( self.__myGroupIds ) == 0:
            return
        
        # Collapse the groups that were organized by category, family, type, 
        # and material.  Also, delete the empty groups to reduce clutter.
        #
        self.sendMessage( 'GROUP_COLLAPSE', ( tuple( self.__myGroupIds ), ) )
        self.sendMessage( 'GROUP_DELETE_EMPTY', () )
        
        self.__myGroupIds = []

    def AUTO_GROUP( self, message ):
        ( collapse, ) = message.data

        self.__myGroupIds = []
        self.__myTransactionId = message.transactionId

        categoriesByModel = {}
        self.__myGroupLabels = {}

        assert( self.__myModels is not None )
        if self.__myModels is None:
            return

        # First group the models' nodes by category, family, type, and material.
        #
        importNodes = CollectNodes( 
            self.__myModels.root, 
            conditions = [lambda node: node.getNodeType() == 'Import' and  \
                                       node.getId().getStageId() == self.__myActiveStageId and \
                                       not node.getIsRemoved()],
            idsOnly=False,
            traversal=kStopOnSuccess
            )

        for importNode in importNodes:
            importNode = CastToOriginalType( importNode )
            importNodeId = importNode.getUniqueId()
        
            modelImport = GetModelImport( importNodeId )
            assert( modelImport is not None )

            modelLeafNodes = CollectModelLeafNodes( modelImport, idsOnly=False )
            
            categories = {}
            uncategorizedNodes = []
            for leafNode in modelLeafNodes:
                # We cannot change nodes that are removed or in a collapsed 
                # hierarchy.
                #
                if leafNode.getIsRemoved() or leafNode.getCollapseRoot() is not None:
                    continue
                    
                objectModelInfo = self.__getRevitObjectModelInfo( leafNode )
                appearance = leafNode.getAppearance()
                material = appearance.getMaterial()
                assert( material )
                materialName = material.getLabel()

                if objectModelInfo is not None:
                    ( category, family, type ) = objectModelInfo
                    
                    if category not in categories:
                        categories[category] = {}
                    
                    if family not in categories[category]:
                        categories[category][family] = {}

                    if type not in categories[category][family]:
                        categories[category][family][type] = {}
                        
                    if materialName not in categories[category][family][type]:
                        categories[category][family][type][materialName] = []                   
                    
                    categories[category][family][type][materialName].append( leafNode.getUniqueId() )

            if len( categories ) > 0:
                categoriesByModel[importNodeId] = categories
        
        # For each grouping by category, family, type, and material, create a 
        # common ancestor (if one doesn't already exist).
        #
        grouped = False
        for importNodeId in AlphanumericSort( categoriesByModel ):
            categories = categoriesByModel[importNodeId]

            for categoryName in AlphanumericSort( categories ):
                category = categories[categoryName]

                for familyName in AlphanumericSort( category ):
                    family = category[familyName]
                    
                    for typeName in AlphanumericSort( family ):
                        type = family[typeName]
                        
                        for materialName in AlphanumericSort( type ):
                            nodes = type[materialName]
                            groupLabel = categoryName + ', ' + familyName + ', ' + typeName + ', ' + materialName
                            grouped = self.__groupNodes( nodes, groupLabel ) or grouped

        # If something was grouped and collapse is requested, do it now.
        #
        if grouped and collapse:
            self.sendMessage( 'AUTO_GROUP_COLLAPSE', () )

## --------------------------------------------------------------------

def instantiate():
    return AutoGroupCustom()


def info():
    customInfo = CustomInfo()
    customInfo.vendor = 'Autodesk'
    customInfo.version = '2.1'
    customInfo.api = '2013'
    customInfo.shortInfo = "An example to group objects by some property."
    customInfo.longInfo = \
"""Define and use a message that groups objects based on some criteria. In this example, the \
objects that have meta-data coming from Revit scenes (families, categories, type, etc.) are \
grouped based on those values.
"""
    return customInfo