Joint Aim (python)

from maya import cmds

import AtomsMath
import AtomsCore
import Atoms
import AtomsUtils
from AtomsMaya.utils import agentgroup
from AtomsMaya.hostbridge.atomsgroup import MayaAtomsGroupHostBridge
 
from Atoms.ui.utils.qthandler import QtWidgets
 
import AtomsMath
import AtomsCore
import Atoms
import AtomsUtils
 
from Atoms.ui.factory import MODULE_WIDGETS_FACTORY
 
from Atoms.ui.components.modules import BaseModuleWidget
 
 
class MyJointAimWidget(BaseModuleWidget):
    def __init__(self, host_bridge, module_name, attributes_map, attribute_properties, parent):
        super(MyJointAimWidget, self).__init__(host_bridge, module_name,
                                             attributes_map, attribute_properties, parent)
 
        self._axis_combo = QtWidgets.QComboBox()
        self._axis_combo.addItems(["X", "Y", "Z"])
        self._axis_combo.setCurrentIndex(attributes_map['axis'].value())
        self.add_widget("Aim Axis", self._axis_combo) 
        self._axis_combo.currentIndexChanged.connect(self._axis_changed)
        
        self._up_axis_combo = QtWidgets.QComboBox()
        self._up_axis_combo.addItems(["X", "Y", "Z"])
        self._up_axis_combo.setCurrentIndex(attributes_map['upAxis'].value())
        self.add_widget("Up Axis", self._up_axis_combo)
        self._up_axis_combo.currentIndexChanged.connect(self._up_axis_changed)
 
    def _axis_changed(self):
        cmds.setAttr("%s.atoms_%s_axis" % (self.app_obj, self.module_name),
                     self._axis_combo.currentIndex())
                     
    def _up_axis_changed(self):
        cmds.setAttr("%s.atoms_%s_upAxis" % (self.app_obj, self.module_name),
                     self._up_axis_combo.currentIndex())
 
 
class MyJointAimModule(Atoms.BehaviourModule):
 
    def __init__(self):
        Atoms.BehaviourModule.__init__(self)
 
        metadata = self.attributes()
        metadata.addEntry("joint", AtomsCore.StringMetadata(""))
        metadata.addEntry("axis", AtomsCore.IntMetadata(0))
        metadata.addEntry("upAxis", AtomsCore.IntMetadata(0))
        metadata.addEntry("upAxisVector", AtomsCore.Vector3Metadata(AtomsMath.V3d(0, 1, 0)))
        metadata.addEntry("target", AtomsCore.Vector3Metadata())
        
        excl_met = AtomsCore.StringArrayMetadata()
        excl_met.set(["axis", "upAxis"])
        
        attr_props = self.attributeProperties()
        mod_props = AtomsCore.MapMetadata()
        mod_props.addEntry(Atoms.GLOBAL_NAMES.BEHAVIOUR._MODULE_.EXCLUDED_METADATAS_AUTOMATIC_BUILD, excl_met)        
        attr_props.addEntry(Atoms.GLOBAL_NAMES.BEHAVIOUR.MODULE_PROPERTIES, mod_props)
 
 
    def endFrame(self, agents, agroup):
        axis_id = self.attributes().getEntry("axis").get()
        up_axis_id = self.attributes().getEntry("upAxis").get()
        target = self.attributes().getEntry("target").get()
        up_vec = self.attributes().getEntry("upAxisVector").get()
        joint = self.attributes().getEntry("joint").value()
 
        if not joint:
            AtomsUtils.Logger.warning("Please provide a valid joint name")
            return
 
        for agent in agents:
            skeleton = agent.agentType().skeleton()
            joint_id = skeleton.jointId(joint)
 
            if joint_id == -1:
                AtomsUtils.Logger.warning("could not find joint with name " +
                                          joint)
                continue

            pose = agent.pose()
            poser = AtomsCore.Poser(skeleton)
 
            wm = poser.getWorldMatrix(pose, joint_id)
            
            pwm = AtomsMath.M44d(wm)
            pwm[3][0] = 0
            pwm[3][1] = 0
            pwm[3][2] = 0
            
            joint_pos = AtomsMath.V3d(wm[3][0], wm[3][1], wm[3][2])
            joint_to_target = target - joint_pos
            
            mtx_right_vec = joint_to_target.cross(up_vec)
            mtx_right_vec.normalize()

            mtx_up_vec = mtx_right_vec.cross(joint_to_target)
            mtx_up_vec.normalize()            
            
            up_axis = AtomsMath.V3d(wm[up_axis_id][0], wm[up_axis_id][1], wm[up_axis_id][2])
            up_axis.normalize()
            
            up_quat = AtomsMath.Quatd()
            up_quat.setRotation(up_axis, mtx_up_vec)
            out_mat = wm * up_quat.toMatrix44()
            
            axis = AtomsMath.V3d(out_mat[axis_id][0], out_mat[axis_id][1], out_mat[axis_id][2])
            axis.normalize()
             
            quat = AtomsMath.Quatd()
            quat.setRotation(axis, joint_to_target)
            
            out_quat = quat * up_quat
            
            pwm = pwm * out_quat.toMatrix44()
            
            out_mat = AtomsMath.M44d()      
            out_mat.translate(joint_pos)
            out_mat = pwm * out_mat
            out_mat.scale(AtomsMath.V3d(1,1,1))

            poser.setWorldMatrix(pose, out_mat, joint_id)
 
 
'''
MAYA SETUP 
'''
def register():
    Atoms.BehaviourModules.instance().registerBehaviourModule("myJointAim",
                                                              MyJointAimModule,
                                                              True)
 
    MODULE_WIDGETS_FACTORY.register('myJointAim', MyJointAimWidget)
 
 


def setup():
    cmds.file(new=True, f=True)
    cmds.tcAtoms(init=True)
    register()
 
    agent_group = agentgroup.create_agent_group()
    ag = MayaAtomsGroupHostBridge()
    ag.set_app_obj(agent_group)
    ag.add_module("gridLayout")
    ag.add_module("stateMachine")
    ag.add_module("myJointAim")
    ag.set_display_type(2)
    
    int_type_str = AtomsCore.IntMetadata.staticTypeStr()
    str_type_str = AtomsCore.StringMetadata.staticTypeStr()
    
    cmds.playbackOptions(min=1, max=50)
    
    s = cmds.spaceLocator()[0]
    cmds.setAttr(s + ".t", 400, 50, -120)
    cmds.setKeyframe(s)
    cmds.currentTime(50)
    cmds.setAttr(s + ".t", 820, 50, 100)
    cmds.setKeyframe(s)
    cmds.currentTime(1)

    ag.set_metadata_value("stateMachine", "state", int_type_str, 1)

    cmds.connectAttr(s + ".t", agent_group + ".atoms_myJointAim_target")
    
    ag.set_metadata_value("myJointAim", "axis", int_type_str, 2)
    ag.set_metadata_value("myJointAim", "joint", str_type_str, "Head")
    
setup()