"""
Ragdoll - Python module for working with ragdolls and MayaBullet.

"""
# Maya
import maya
maya.utils.loadStringResourcesForModule(__name__)

import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
from pymel.core import *
import maya.cmds
# Python standard library
import math
import logging
from collections import deque
# MayaBullet
import maya.app.mayabullet.RigidBody as RigidBody
import maya.app.mayabullet.RigidBodyConstraint as RigidBodyConstraint
import maya.app.mayabullet.CommandWithOptionVars as CommandWithOptionVars

from maya.app.mayabullet.RigidBodyConstraint import eConstraintType, eConstraintLimitType
from maya.app.mayabullet.RigidBody import eBodyType, eShapeType, eAxisType

from maya.app.mayabullet import logger as logger


################################### GLOBALS ####################################

RB_XFORM_PREFIX = "jointToRigidBody"

ROTATE_MIN = -360
ROTATE_MAX =  360

# Angular constraint defaults
DEFAULT_ANGULAR_DAMPING     = 0.5
DEFAULT_ANGULAR_SOFTNESS    = 0.0
DEFAULT_ANGULAR_RESTITUTION = 0.0
# Capsule defaults
DEFAULT_CAPSULE_LENGTH      = 0.8  # as a proportion of bone length
DEFAULT_CAPSULE_RADIUS      = 0.1  # as a proportion of bone length
DEFAULT_CAPSULE_MASS        = 1.0
DEFAULT_NAME_SEPARATOR      = '_'


############################### HELPER FUNCTIONS ###############################

def _calcCapsuleRadius( jointStart, jointEnd, 
                        radiusRatio=DEFAULT_CAPSULE_RADIUS ):
    """
    Given a bone (specified by the joints on either end), return the
    radius of the capsule for this bone.
    """
    # This is a simple placeholder implementation. Ideas for more
    # sophisticated solutions:
    #
    # TODO: implement based on bounds around static mesh shapes
    # TODO: implement based on clusters for skinned meshes
    return jointEnd.getTranslation().length() * radiusRatio
# end


def _getNodesToVisit( rootJointName ):
    """
    Returns a list of PyNodes for traversal. If the rootJointName is
    None, uses the current selection.
    """
    # figure out where the traversal will start
    if ( rootJointName is None ):
        nodesToVisit = ls( long=True, selection=True, type='joint' )
    else:
        try:
            nodesToVisit = [ PyNode(rootJointName) ]
            assert( isinstance(nodesToVisit[0], nodetypes.Joint) )
        except:
            logger.error( maya.stringTable['y_Ragdoll.kUnableStartJoint'] % rootJointName ) 
            raise
        # end-try
    # end-if    

    return nodesToVisit
# end


def _getCapsule( startJoint, endJoint ):
    """
    Given a parent (start) joint and child (end) joint, retrieve the
    (user editable) rigid body transform and the rigid body shape that
    is associated with this bone.

    Returns (rbXform, rbShape) on success, or (None, None) on failure.
    """
    if ( startJoint is None or endJoint is None ):
        return (None, None)

    childXforms = startJoint.getChildren( type="transform" )
    for childXform in childXforms:
        if ( childXform.getTranslation() == endJoint.getTranslation() * 0.5 ):
            # NOTE: assumes only the bullet transform is the only
            #       child of the capsule rb transform.
            # NOTE: assumes there is only one rigid body under the
            #       bullet transform
            return ( childXform,
                     childXform.childAtIndex(0)\
                         .getChildren( type="bulletRigidBodyShape" )[0] )
        # end-if
    # end-for
    
    return (None, None)
# end


def _createCapsule( currentJoint, childJoint, bAttachToJoint=True,
                    bodyType=eBodyType.kKinematicRigidBody, mass=0.0,
                    boneLengthRatio   = DEFAULT_CAPSULE_LENGTH,
                    lengthRadiusRatio = DEFAULT_CAPSULE_RADIUS,
                    transformName=None ):
    """
    Create a capsule collider around the bone from currentJoint to
    childJoint. If bAttachToJoint is True, the capsule will be
    parented to the current joint, otherwise it will be created at the
    top level of the scene. BodyType and mass can be specified, and
    transform name can be used to specify the name of the new
    top-level rigid body transform.

    This implementation is based on the current rigid body hierarchy,
    which looks like:

    rbXform
      --> bulletXform
            --> rbShape (capsule)

    Returns (rbXform, rbShape)
    """
    if ( transformName is None ):
        transformName = '%s#' % RB_XFORM_PREFIX

    if ( bAttachToJoint ): 
        parent=currentJoint
    else: 
        parent=None

    rbXform = createNode( 'transform', name=transformName,
                          parent=parent, skipSelect=True )

    # these values will work when adding capsules to a joint hierarchy
    boneVector   = childJoint.getTranslation()
    center  =  datatypes.Vector( boneVector.length() * .5, 0 ,0 )
    worldOffset  = datatypes.Vector( 0, 0, 0 )

    if ( not bAttachToJoint ):
        # if this capsule won't be inheriting the transform from the 
        # parent joint, these values need to be adjusted for world-space
        worldOffset = currentJoint.getTranslation( space="world" )
        boneVector  = childJoint.getTranslation(space="world") \
                          - currentJoint.getTranslation(space="world")
        center =  datatypes.Vector( boneVector.length() * .5, 0 ,0 )  # rotation is from the x axis.
    # end-if

    rotateAngles = datatypes.Quaternion( datatypes.Vector(1,0,0), 
                                         boneVector ).asEulerRotation()
    # a workaround for some PyMEL bugs handling rotations...
    cmds.rotate( math.degrees(rotateAngles.x),
                 math.degrees(rotateAngles.y),
                 math.degrees(rotateAngles.z),
                 rbXform.longName() )
    rbXform.setTranslation( worldOffset)        

    logger.info( maya.stringTable[ 'y_Ragdoll.kAddingRB']   % rbXform )
    (rbParentName, rbName) \
        = RigidBody.CreateRigidBody \
              .command( transformName=rbXform.longName(),
                        colliderShapeType=eShapeType.kColliderCapsule,
                        axis=eAxisType.kXAxis,
                        ignoreShape=True,
                        bAttachSelected=False )
    rbNode = PyNode( rbName )
    rbNode.bodyType.set( bodyType )
    rbNode.mass.set( mass )
    rbNode.centerOfMass.set( center )  
    capsuleRadius = _calcCapsuleRadius( currentJoint, childJoint, 
                                        lengthRadiusRatio )
    rbNode.length \
        .set( max(0.0, 
                  math.fabs(boneVector.length() * boneLengthRatio) 
                      - capsuleRadius) )
    rbNode.radius.set( capsuleRadius )

    return ( rbXform, rbNode )
# end


def _getRagdollCapsule( parentJoint, currentJoint, 
                        jointNameSeparator=DEFAULT_NAME_SEPARATOR ):
    """
    Retrieves the rigid body from the ragdoll that corresponds to the
    give pair of joints. This method relies on the naming convention
    used for creating the transforms for the rigid bodies.

    Returns (rbXform, rb), or (None, None) on failure.
    """
    if ( parentJoint is None or currentJoint is None ):
        return (None, None)

    capsuleXform \
        = PyNode( "%s%s%s" % (parentJoint, jointNameSeparator, currentJoint) )
    return ( capsuleXform, 
             # NOTE: assumes there is only one transform under the capsule
             capsuleXform.getChildren( type="bulletRigidBodyShape" )[0] )
# end


def _applyJointLimits( joint, constraint ):
    """
    joint - PyNode of the joint to get the limits from
    constraint - PyNode of the constraint shape for a 6-DOF constraint
                 to apply the limits to 
    """
    minRotLimitEnable = joint.minRotLimitEnable.get()
    maxRotLimitEnable = joint.maxRotLimitEnable.get()
    minRotLimit = joint.minRotLimit.get()
    maxRotLimit = joint.maxRotLimit.get()

    if ( minRotLimitEnable[0] or maxRotLimitEnable[0] ):
        if ( minRotLimit[0] == maxRotLimit[0] ):
            constraint.angularConstraintX \
                .set( eConstraintLimitType.kRBConstraintLimitLocked )
        else:
            constraint.angularConstraintX \
                .set( eConstraintLimitType.kRBConstraintLimitLimited )
            # Since Bullet doesn't make the distinction between
            # enabling the min or max of a limit, if the joint has one
            # end of the range is disabled, set the min/max angle to
            # be effectively disabled in Bullet.
            if ( not minRotLimitEnable[0] ):
                minRotLimit[0] = ROTATE_MIN
            if ( not maxRotLimitEnable[0] ):
                maxRotLimit[0] = ROTATE_MAX
        # end-if

    if ( minRotLimitEnable[1] or maxRotLimitEnable[1] ):
        if ( minRotLimit[1] == maxRotLimit[1] ):
            constraint.angularConstraintY \
                .set( eConstraintLimitType.kRBConstraintLimitLocked )
        else:
            constraint.angularConstraintY \
                .set( eConstraintLimitType.kRBConstraintLimitLimited )
            if ( not minRotLimitEnable[1] ):
                minRotLimit[1] = ROTATE_MIN
            if ( not maxRotLimitEnable[1] ):
                maxRotLimit[1] = ROTATE_MAX
        # end-if

    if ( minRotLimitEnable[2] or maxRotLimitEnable[2] ):
        if ( minRotLimit[2] == maxRotLimit[2] ):
            constraint.angularConstraintZ \
                .set( eConstraintLimitType.kRBConstraintLimitLocked )
        else:
            constraint.angularConstraintZ \
                .set( eConstraintLimitType.kRBConstraintLimitLimited )
            if ( not minRotLimitEnable[2] ):
                minRotLimit[2] = ROTATE_MIN
            if ( not maxRotLimitEnable[2] ):
                maxRotLimit[2] = ROTATE_MAX
        # end-if

    # push the limits through to the constraint regardless of what's
    # enabled, since they should be harmless for any disabled axes.
    constraint.angularConstraintMin.set( minRotLimit ) 
    constraint.angularConstraintMax.set( maxRotLimit )
# end


def _addWidget( widgetFcn, optionVarName, widgetDict, **kwargs ):
    """
    Helper for adding UI widgets in addOptionDialogWidget() callbacks.
    """
    widget = widgetFcn( **kwargs )
    widgetDict[optionVarName] = ( widgetFcn, widget )
# end


def _addCapsuleWidgets( widgetDict, bIncludeMass=True ):
    """
    Adds the widgets for editing capsule attributes.
    """
    if ( bIncludeMass ):
        _addWidget( maya.cmds.floatSliderGrp, 'capsuleMass', widgetDict,
                    label=maya.stringTable[ 'y_Ragdoll.kCapsuleMass']  ,
                    minValue=0, maxValue=100, fieldMaxValue=9999999 )
    # end-if
    _addWidget( maya.cmds.floatSliderGrp, 'capsuleBoneRatio', widgetDict,
                label=maya.stringTable[ 'y_Ragdoll.kCapsuleBoneLengthRatio']  ,
                minValue=0, maxValue=1, fieldMaxValue=9999999 )
    _addWidget( maya.cmds.floatSliderGrp, 'capsuleRadiusRatio', widgetDict,
                label=maya.stringTable[ 'y_Ragdoll.kCapsuleRadiusLengthRatio']  ,
                minValue=0, maxValue=1, fieldMaxValue=9999999 )
# end


############################### PUBLIC API #####################################

def addCapsulesToSkeleton( rootJointName      = None, 
                           capsuleBoneRatio   = DEFAULT_CAPSULE_LENGTH,
                           capsuleRadiusRatio = DEFAULT_CAPSULE_RADIUS ):
    """
    This method traverses a joint hierarchy, adding kinematic rigid
    body capsules to the bones. If the name of the root joint is
    provided, the traversal will begin at that joint. Otherwise, the
    currently selected joint(s) will be used.
    """
    initialSelection = ls( long=True, selection=True )
    nodesToVisit = deque( _getNodesToVisit( rootJointName ) )

    while ( len(nodesToVisit) > 0 ):
        currentJoint = nodesToVisit.popleft()
        logger.debug( maya.stringTable[ 'y_Ragdoll.kVisiting']   % currentJoint )

        # within the current iteration, we will consider only other
        # joints attached to the current joint
        children = [ child for child in currentJoint.getChildren() \
                         if isinstance(child, nodetypes.Joint) ]
        # future iterations need only consider joints with children
        nodesToVisit.extend( [child for child in children \
                                  if child.numChildren() > 0] )
        prevCapsules = [child for child in currentJoint.getChildren() \
                            if child.shortName().startswith(RB_XFORM_PREFIX)]

        # examine the bones between this joint and its child joints...
        for childJoint in children:
            logger.debug( maya.stringTable[ 'y_Ragdoll.kBoneFromTo']   % (currentJoint, childJoint) )
            boneVector = childJoint.getTranslation()

            bCapsuleExists = False
            for prevCapsuleXform in prevCapsules:
                if prevCapsuleXform.getTranslation() == boneVector * 0.5:
                    bCapsuleExists = True
                    break
                # end-if
            # end-for
            if ( bCapsuleExists ): 
                logger.warn( maya.stringTable[ 'y_Ragdoll.kSkippingBone']   )
                continue
            # end-if

            _createCapsule( currentJoint, childJoint, bAttachToJoint=True,
                            # colliders for animation should be kinematic, 
                            # with zero mass
                            bodyType=eBodyType.kKinematicRigidBody, mass=0.0,
                            boneLengthRatio=capsuleBoneRatio,
                            lengthRadiusRatio=capsuleRadiusRatio )
        # end-for
    # end-while

    # restore the initial selection
    select( initialSelection, replace=True )

# end addCapsulesToSkeleton()


class AddColliders( CommandWithOptionVars.CommandWithOptionVars ):
    """
    OptionBox wrapper for addCapsulesToSkeleton()
    """

    def __init__( self ):
        super( AddColliders, self ).__init__()

        self.commandName       = 'Add Colliders'
        self.l10nCommandName   = maya.stringTable['y_Ragdoll.kAddColliders'] 
        self.optionVarPrefix   = 'bullet_Ragdoll_'
        self.optionVarDefaults = {
            'capsuleBoneRatio'    : DEFAULT_CAPSULE_LENGTH,
            'capsuleRadiusRatio'  : DEFAULT_CAPSULE_RADIUS,
        }
    # end


    @staticmethod
    def command( **kwargs ):
        """
        """
        # If command echoing is off, echo this short line.
        if (not maya.cmds.commandEcho(query=True, state=True)):
            print("Ragdoll.AddColliders()..executeCommandCB()")
        addCapsulesToSkeleton( **kwargs )
    # end


    def addOptionDialogWidgets( self ):
        """
        """
        # dict format: {optionVarDictKey, (widgetClass, widget)}
        widgetDict = {} 

        separator( style="none", height=5 )
        _addCapsuleWidgets( widgetDict, bIncludeMass=False )

        return widgetDict
    # end

# end class        


def createRagdoll( rootJoint=None,
                   angularDamping     = DEFAULT_ANGULAR_DAMPING,
                   angularSoftness    = DEFAULT_ANGULAR_SOFTNESS,
                   angularRestitution = DEFAULT_ANGULAR_RESTITUTION,
                   capsuleBoneRatio   = DEFAULT_CAPSULE_LENGTH,
                   capsuleRadiusRatio = DEFAULT_CAPSULE_RADIUS,
                   capsuleMass        = DEFAULT_CAPSULE_MASS,
                   jointNameSeparator = DEFAULT_NAME_SEPARATOR ):
    """
    Creates a ragdoll of capsules joined by constraints, that matches
    the skeleton starting from the joint named by rootJoint. If no
    root is specified, the current selection is used.
    """
    ragdollRootNode = createNode( 'transform', name="Ragdoll#",
                                   skipSelect=True )
                          
    nodesToVisit = deque( _getNodesToVisit(rootJoint) )
    currentJoint = None
    while ( len(nodesToVisit) > 0 ):
        logger.debug( maya.stringTable[ 'y_Ragdoll.kStartingIterCurrJoint']   \
                           % (currentJoint, nodesToVisit) )
        currentJoint = nodesToVisit.popleft()
        # NOTE: assumes joints are directly connected to each other
        try:
            parentJoint = currentJoint.firstParent()
            if ( not isinstance(parentJoint, nodetypes.Joint) ):
                parentJoint = None
        except:
            parentJoint = None
        # end-try

        logger.debug( maya.stringTable[ 'y_Ragdoll.kVisiting']   % currentJoint )
        (rbAXform, rbA) = _getRagdollCapsule( parentJoint, currentJoint,
                                              jointNameSeparator )
        childJoints = currentJoint.getChildren( type="joint" )
        nodesToVisit.extend( childJoints )

        prevChildJoint = None
        for childJoint in childJoints:
            # Here's what we're working with within this loop:
            #
            # parent           current           child
            #  joint --[rbA]--> joint --[rbB]--> joint
            (rbBXform, rbB) \
                = _createCapsule( currentJoint, childJoint, False,
                                  eBodyType.kDynamicRigidBody, capsuleMass,
                                  capsuleBoneRatio, capsuleRadiusRatio,
                                  "%s%s%s" % (currentJoint.shortName(),
                                              jointNameSeparator,
                                              childJoint.shortName()) )
            rbBXform.setParent( ragdollRootNode )
            rbBName = None
            if ( rbB is not None ): rbBName = rbB.longName()
            rbBXformName = None
            if ( rbBXform is not None ): rbBXformName = rbBXform.longName()

            # figure out what to anchor this rigid body to
            if ( rbA is None and prevChildJoint is not None ): 
                # prevChildJoint implies that a sibling capsule was
                # created in a previous iteration, which this capsule
                # can be constrained to. 
                (sibXform, sib) \
                    = _getRagdollCapsule( currentJoint, prevChildJoint,
                                          jointNameSeparator )
                rbAnchorName = sib.longName()
            elif ( rbA is not None ):
                rbAnchorName = rbA.longName()
            else:
                # NOTE: don't create constraints without an anchor, or
                # else the ragdoll will end up constrained to the world.
                rbAnchorName = None
            # end-if

            if ( rbAnchorName is not None ):
                logger.info( maya.stringTable[ 'y_Ragdoll.kCreatingConstr']   \
                                  % (rbAnchorName, rbB) )
                # "constraint" abbreviated "constr"
                constrName \
                    = RigidBodyConstraint.CreateRigidBodyConstraint \
                          .command( constraintType=eConstraintType.kRBConstraintSixDOF,
                                    rigidBodyA=rbAnchorName, 
                                    rigidBodyB=rbBName, 
                                    parent=None )[0]
                constrNode = PyNode( constrName )
                constrXformNode = constrNode.firstParent()

                # configure the transform
                rename( constrXformNode, 
                        "constraint_%s" % currentJoint.shortName() )
                constrXformNode.setParent( ragdollRootNode )
                constrXformNode.setTranslation( 
                    currentJoint.getTranslation(space="world") )
                constrXformNode.setRotation(
                    currentJoint.getRotation(space="world") )

                # configure the constraint
                constrNode.angularDamping.set( angularDamping )
                constrNode.angularSoftness.set( angularSoftness )
                constrNode.angularRestitution.set( angularRestitution )

                # lock the linear motion, to behave like a point constraint
                constrNode.linearConstraintX \
                    .set( eConstraintLimitType.kRBConstraintLimitLocked )
                constrNode.linearConstraintY \
                    .set( eConstraintLimitType.kRBConstraintLimitLocked )
                constrNode.linearConstraintZ \
                    .set( eConstraintLimitType.kRBConstraintLimitLocked )

                # set the rotational limits to match the joint
                _applyJointLimits( currentJoint, constrNode )
            # end-if

            prevChildJoint = childJoint
        # end-for
    # end-while

    select( ragdollRootNode, replace=True )
# end createRagdoll()


class CreateRagdoll( CommandWithOptionVars.CommandWithOptionVars ):
    """
    OptionBox wrapper for createRagdoll().
    """

    def __init__( self ):
        super( CreateRagdoll, self ).__init__()

        self.commandName       = 'Create Ragdoll'
        self.l10nCommandName   = maya.stringTable['y_Ragdoll.kCreateRagdoll'] 
        self.optionVarPrefix   = 'bullet_Ragdoll_'
        self.optionVarDefaults = {
            'angularDamping'      : DEFAULT_ANGULAR_DAMPING,
            'angularSoftness'     : DEFAULT_ANGULAR_SOFTNESS,
            'angularRestitution'  : DEFAULT_ANGULAR_RESTITUTION,
            'capsuleMass'         : DEFAULT_CAPSULE_MASS,
            'capsuleBoneRatio'    : DEFAULT_CAPSULE_LENGTH,
            'capsuleRadiusRatio'  : DEFAULT_CAPSULE_RADIUS,
            'jointNameSeparator'  : DEFAULT_NAME_SEPARATOR,
        }
    # end


    @staticmethod
    def command( **kwargs ):
        """
        Callback to execute the command.
        """
        
        # If command echoing is off, echo this short line.
        if (not maya.cmds.commandEcho(query=True, state=True)):
            print("Ragdoll.CreateRagdoll().executeCommandCB()")
        createRagdoll( **kwargs )
    # end


    def addOptionDialogWidgets( self ):
        """
        Callback to set up the OptionBox UI.
        """
        # dict format: {optionVarDictKey, (widgetClass, widget)}
        widgetDict = {} 

        separator( style="none", height=5 )
        _addCapsuleWidgets( widgetDict )

        separator( style="none", height=5 )
        _addWidget( maya.cmds.floatSliderGrp, 'angularDamping', widgetDict,
                    label=maya.stringTable[ 'y_Ragdoll.kJointAngDampLabel']  , 
                    minValue=0, maxValue=1, fieldMaxValue=9999999 )
        _addWidget( maya.cmds.floatSliderGrp, 'angularSoftness', widgetDict,
                    label=maya.stringTable[ 'y_Ragdoll.kJointAngSoftLabel']  , 
                    minValue=0, maxValue=1, fieldMaxValue=9999999 )
        _addWidget( maya.cmds.floatSliderGrp, 'angularRestitution', widgetDict,
                    label=maya.stringTable[ 'y_Ragdoll.kJointAngRestitutionLabel']  , 
                    minValue=0, maxValue=1, fieldMaxValue=9999999 )

        _addWidget( maya.cmds.textFieldGrp, 'jointNameSeparator', widgetDict,
                    label=maya.stringTable[ 'y_Ragdoll.kJointNameSepLabel']   )

        return widgetDict
    # end

# end class        
# Copyright (C) 1997-2011 Autodesk, Inc., and/or its licensors.
# All rights reserved.
#
# The coded instructions, statements, computer programs, and/or related
# material (collectively the "Data") in these files contain unpublished
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its licensors,
# which is protected by U.S. and Canadian federal copyright law and by
# international treaties.
#
# The Data is provided for use exclusively by You. You have the right to use,
# modify, and incorporate this Data into other products for purposes authorized 
# by the Autodesk software license agreement, without fee.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. AUTODESK
# DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED WARRANTIES
# INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF NON-INFRINGEMENT,
# MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE, OR ARISING FROM A COURSE 
# OF DEALING, USAGE, OR TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS
# LICENSORS BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL,
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK AND/OR ITS
# LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY OR PROBABILITY OF SUCH DAMAGES.

