'''
Pose Tool

Author: Jason Hayes (jason.hayes@rockstarsandiego.com)

Description:  Based on the 3dsmax Pose2Pose tool, allows the storing, loading and blending of poses.

Note: I didn't know MotionBuilder (as in I loaded up MoBu for the first time in MY LIFE) at the time of writing this tool, 
so future apologies to whomever sees my code and is cursing to themselves "WTF was that dude thinking writing it like this?!".
'''

import os
import xml.etree.cElementTree

from pyfbsdk import *


## Misc ##

def getTabs( numTabs ):
    ''' 
    Returns a string with the specified number of tabs.  Used when creating the scene xml for the pose list. 
    '''
    tabs = ''
    
    for i in range( 0, numTabs ):
        tabs += '\t'
        
    return tabs
        
        
## Classes ##
       
class RsPoseNode( object ):
    ''' 
    Represents a node in the current scene, but does not contain a direct pointer to the scene object.
    Call getSceneNode() to get a pointer to the real object in the scene.
    '''  

    def __init__( self ):
        
        # The name of the scene node.
        self.name       = None
        
        # Saved pose data.
        self.position   = None
        self.rotation   = None
        self.scale      = None
        
        # Snapshot current node transform.
        self.__snapshotPosition = None
        self.__snapshotRotation = None
        self.__snapshotScale    = None
        
            
    def __interpolateVectors( self, vecA, vecB, time ):
        '''
        Performs a linear interpolation between two vectors and a time constant.
        
        Pissed that I can't do multiplication on an FBVector3d object, which makes this way more code
        than is needed.
        '''
        aX = vecA.GetList()[ 0 ]
        aY = vecA.GetList()[ 1 ]
        aZ = vecA.GetList()[ 2 ]
        
        bX = vecB.GetList()[ 0 ]
        bY = vecB.GetList()[ 1 ]
        bZ = vecB.GetList()[ 2 ]
        
        x = bX * ( 1.0 - time ) + aX * time
        y = bY * ( 1.0 - time ) + aY * time
        z = bZ * ( 1.0 - time ) + aZ * time
        
        return FBVector3d( x, y, z )
        
            
    def snapshot( self ):
        '''
        Creates a snapshot of the current node transform.
        '''
        sceneNode = self.getSceneNode()
        
        if sceneNode:
            self.__snapshotPosition = FBVector3d( sceneNode.Translation.Data )
            self.__snapshotRotation = FBVector3d( sceneNode.Rotation.Data )
            self.__snapshotScale = FBVector3d( sceneNode.Scaling.Data )
            
        else:
            assert 0, "Could not locate object ({0}) in the scene!  The pose tool data has somehow gotten out of sync, or the character is not compatible with this file.".format( self.name )
            
            
    def blend( self, pct ):
        '''
        Blends the current node to the stored pose transform, based on the supplied percentage. The
        percentage should be normalized down to 0 - 1.
        '''
        sceneNode = self.getSceneNode()
        
        if sceneNode:    
            sceneNode.Translation = self.__interpolateVectors( self.position, self.__snapshotPosition, pct )
            sceneNode.Rotation = self.__interpolateVectors( self.rotation, self.__snapshotRotation, pct )
            sceneNode.Scaling = self.__interpolateVectors( self.scale, self.__snapshotScale, pct )
            
                
    def key( self ):
        '''
        Sets a key for the pose at the current frame time.
        '''
        sceneNode = self.getSceneNode()
        
        if sceneNode:
            sceneNode.Translation.Key()
            sceneNode.Rotation.Key()
            sceneNode.Scaling.Key()
            
        
    def getSceneNode( self ):
        '''
        Returns the real scene node, if found.
        '''
        return FBFindModelByLabelName( self.name )
        
        
    def save( self, tabLevel ):
        '''
        Creates an xml representation of the object as a string and returns it.
        '''
        xml = '{0}<node name="{1}">\n'.format( getTabs( tabLevel ), self.name )
        tabLevel += 1
        
        xPos, yPos, zPos = self.position.GetList()
        xRot, yRot, zRot = self.rotation.GetList()
        xScale, yScale, zScale = self.scale.GetList()
        
        xml += '{0}<position>{1} {2} {3}</position>\n'.format( getTabs( tabLevel ), xPos, yPos, zPos )
        xml += '{0}<rotation>{1} {2} {3}</rotation>\n'.format( getTabs( tabLevel ), xRot, yRot, zRot )
        xml += '{0}<scale>{1} {2} {3}</scale>\n'.format( getTabs( tabLevel ), xScale, yScale, zScale )
        
        tabLevel -= 1
        xml += '{0}</node>\n'.format( getTabs( tabLevel ) )
        
        return xml
        
    
class RsPose( object ):
    '''
    Represents a stored pose.
    '''
    def __init__( self ):
        
        # The name of the pose.
        self.name       = None
        
        # The RsPoseNode objects that represent the pose.
        self.nodes      = []
        
        
    def save( self, tabLevel ):
        '''
        Creates an xml representation of the pose as a string and returns it.
        
        tabLevel: The number of tabs (indentation) where the xml block should start.
        '''
        xml = '{0}<pose name="{1}">\n'.format( getTabs( tabLevel ), self.name )
        tabLevel += 1
        
        xml += '{0}<nodes>\n'.format( getTabs( tabLevel ) )
        tabLevel += 1
        
        for node in self.nodes:
            xml += node.save( tabLevel )
        
        tabLevel -= 1
        xml += '{0}</nodes>\n'.format( getTabs( tabLevel ) )
        
        tabLevel -= 1
        xml += '{0}</pose>\n'.format( getTabs( tabLevel ) )
        
        return xml
        
         
      
class RsPoseManager( object ):
    '''
    Static interface for managing the poses.
    '''
    
    # Store a version number in case of xml format changes.
    version = 1.0
    
    # Track if the poses need to be saved.
    dirty = False
    
    # Dictionary of all the loaded poses.
    # key: The pose name.
    # value: RsPose object.
    poses = {}
    
    # File extension for the saved pose files.
    fileExtension = 'p2px'
    
    # Store the last loaded pose file.
    lastPoseFile = None
    
    namespace = None
        
  
    @staticmethod
    def getModelSelection():
        '''
        Return the current model selection.
        '''
        modelList = FBModelList()
        FBGetSelectedModels( modelList )
        
        return modelList
     
     
    @staticmethod
    def selectPoseControls( poseName ):
        if poseName in RsPoseManager.poses:
            pose = RsPoseManager.poses[ poseName ]
            
            # Clear current selection.
            modelList = FBModelList()
            FBGetSelectedModels( modelList, None, True )
            
            for model in modelList:
                model.Selected = False
            
            # Now select just the nodes for this pose.
            for node in pose.nodes:
                obj = node.getSceneNode()
                
                if obj:
                    obj.Selected = True
     
    
    @staticmethod   
    def clearAll( ):
        '''
        Clear all poses.
        '''
        RsPoseManager.dirty = True
        RsPoseManager.poses = {}
        
        
    @staticmethod
    def snapshot( poseName ):
        '''
        Snapshot deltas for the supplied pose.
        '''
        if poseName in RsPoseManager.poses:
            pose = RsPoseManager.poses[ poseName ]
            
            for node in pose.nodes:
                node.snapshot()
        
        
    @staticmethod
    def blendPose( poseName, pct, key = True ):
        '''
        Blends a pose using the supplied percentage.
        '''
        if poseName in RsPoseManager.poses:
            pose = RsPoseManager.poses[ poseName ]
            
            for node in pose.nodes:
                node.blend( pct )
                
                if key:
                    node.key()
                    
    
    @staticmethod
    def keyPose( poseName ):
        '''
        Sets an animation key for a pose.
        '''
        if poseName in RsPoseManager.poses:
            pose = RsPoseManager.poses[ poseName ]
            
            for node in pose.nodes:
                node.key()
                    
    
    @staticmethod
    def savePoses( filename ):
        '''
        Save the current set of poses to file.
        '''
        poseFile = open( filename, 'w' )

        xml = '<?xml version="1.0" encoding="utf-8"?>\n'
        xml += '<poses version="{0}" namespace="{1}">\n'.format( RsPoseManager.version, RsPoseManager.namespace )
        tabLevel = 1
        
        for poseName, pose in RsPoseManager.poses.iteritems():
            xml += pose.save( tabLevel )
        
        xml += '</poses>\n'
        
        poseFile.write( xml )
        poseFile.close()
        
        RsPoseManager.dirty = False
        
    @staticmethod
    def loadLastPoses( ):
        if RsPoseManager.lastPoseFile != None:
            if os.path.exists( RsPoseManager.lastPoseFile ):
                RsPoseManager.loadPoses( RsPoseManager.lastPoseFile )
  
  
    @staticmethod      
    def loadPoses( filename ):
        '''
        Load a set of poses from file.
        '''
        if os.path.exists( filename ):
            RsPoseManager.poses = {}
            
            doc = xml.etree.cElementTree.parse( filename )
            root = doc.getroot()
            
            if root:
                RsPoseManager.namespace = root.get( 'namespace' )
                #RsPoseManager.version = root.get( 'version' )
                
                xmlPoses = root.findall( 'pose' )
                
                for xmlPose in xmlPoses:
                    pose = RsPose()
                    pose.name = xmlPose.get( 'name' )
                    
                    xmlNodes = xmlPose.find( 'nodes' )
                    
                    for xmlNode in xmlNodes:
                        poseNode = RsPoseNode()
                        
                        poseNode.name = xmlNode.get( 'name' )
                        
                        # Get and assign transform information.
                        xPos, yPos, zPos = str( xmlNode.find( 'position' ).text ).split( ' ' )
                        xRot, yRot, zRot = str( xmlNode.find( 'rotation' ).text ).split( ' ' )
                        xScale, yScale, zScale = str( xmlNode.find( 'scale' ).text ).split( ' ' )
                        
                        poseNode.position = FBVector3d( float( xPos ), float( yPos ), float( zPos ) )
                        poseNode.rotation = FBVector3d( float( xRot ), float( yRot ), float( zRot ) )
                        poseNode.scale = FBVector3d( float( xScale ), float( yScale ), float( zScale ) )
                        
                        pose.nodes.append( poseNode )
                        
                    RsPoseManager.poses[ pose.name ] = pose
                    
                RsPoseManager.lastPoseFile = filename
                
        else:
            msg = FBMessageBox( "Pose Tool", "The post file ({0}) does not exist!".format( filename ), "OK" )
            del( msg )
        
    
    @staticmethod
    def renamePose( oldPoseName, newPoseName ):
        '''
        Rename an existing pose.
        '''
        if oldPoseName in RsPoseManager.poses:
            RsPoseManager.dirty = True
                        
            pose = RsPoseManager.poses[ oldPoseName ]
            pose.name = newPoseName
            
            RsPoseManager.poses.pop( oldPoseName )
            RsPoseManager.poses[ newPoseName ] = pose
            
        else:
            msg = FBMessageBox( "Pose Tool", "The supplied pose name ({0}) does not exist!".format( oldPoseName ), "OK" )
            del( msg ) 
        
    
    @staticmethod
    def deletePose( poseName ):
        '''
        Delete an existing pose.
        '''
        if poseName in RsPoseManager.poses:
            RsPoseManager.dirty = True
            RsPoseManager.poses.pop( poseName )
            
        else:
            msg = FBMessageBox( "Pose Tool", "The supplied pose name ({0}) does not exist!".format( poseName ), "OK" )
            del( msg )

    
    @staticmethod
    def updatePose( poseName ):
        '''
        Update an existing pose to use a different set of nodes.
        '''
        if poseName in RsPoseManager.poses:
            RsPoseManager.dirty = True
                        
            selection = RsPoseManager.getModelSelection()
            pose = RsPoseManager.poses[ poseName ]
            pose.nodes = []
            
            for obj in selection:
                node = RsPoseNode()
                node.name = obj.LongName
                node.position = FBVector3d( obj.Translation.Data )
                node.rotation = FBVector3d( obj.Rotation )
                node.scale = FBVector3d( obj.Scaling )
                
                pose.nodes.append( node )
                
                lNamespace = obj.LongName.partition(":")
                if lNamespace[1] != "":
                    RsPoseManager.namespace = lNamespace[0]
                else:
                    RsPoseManager.namespace = "None"
                
            RsPoseManager.poses[ poseName ] = pose
            
        else:
            msg = FBMessageBox( "Pose Tool", "The supplied pose name ({0}) does not exist!".format( poseName ), "OK" )
            del( msg )
            
    
    @staticmethod
    def createNewPose( poseName ):
        '''
        Create a new pose based on the current selection.
        '''
        selection = RsPoseManager.getModelSelection()
        
        if len( selection ) > 0:
            RsPoseManager.dirty = True
            
            pose = RsPose()
            pose.name = poseName
            
            for obj in selection:
                node = RsPoseNode()
                node.name = obj.LongName
                node.position = FBVector3d( obj.Translation.Data )
                node.rotation = FBVector3d( obj.Rotation )
                node.scale = FBVector3d( obj.Scaling )
                
                pose.nodes.append( node )
                
                lNamespace = obj.LongName.partition(":")
                if lNamespace[1] != "":
                    RsPoseManager.namespace = lNamespace[0]
                else:
                    RsPoseManager.namespace = "None"
                
            RsPoseManager.poses[ poseName ] = pose
            
        else:
            msg = FBMessageBox( "Pose Tool", "No objects selected!", "OK" )
            del( msg )