Joint Aim (python)

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, parent):
        super(MyJointAimWidget, self).__init__(host_bridge, module_name,
                                             attributes_map, 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)

    @staticmethod
    def excluded_metadatas_from_automatic_build():
        return ["axis"]

    def _axis_changed(self):
        cmds.setAttr("%s.atoms_%s_axis" % (self.app_obj, self.module_name),
                     self._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("target", AtomsCore.Vector3Metadata())


    def endFrame(self, agents, agroup):
        axis_id = self.attributes().getEntry("axis").get()
        target = self.attributes().getEntry("target").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)
            axis = AtomsMath.V3d(wm[axis_id][0], wm[axis_id][1], wm[axis_id][2])

            joint_pos = AtomsMath.V3d(wm[3][0], wm[3][1], wm[3][2])
            joint_to_target = target - joint_pos

            quat = AtomsMath.Quatd()
            quat.setRotation(axis, joint_to_target)

            wm = quat.toMatrix44() * wm
            poser.setWorldMatrix(pose, wm, joint_id)


'''
MAYA SETUP 
'''
def register():
    Atoms.BehaviourModules.instance().registerBehaviourModule("myJointAim",
                                                              MyJointAimModule,
                                                              True)

    MODULE_WIDGETS_FACTORY.register('myJointAim', MyJointAimWidget)


def setup():
    cmds.tcAtoms(init=True)
    agn = cmds.createNode('tcAgentGroupNode')
    cmds.setAttr(agn + ".displayType", 1)

    register()

    cmds.setAttr(agn + ".modules[0].moduleType", "gridLayout", type="string")
    cmds.setAttr(agn + ".modules[0].moduleName", "gridLayout", type="string")

    cmds.setAttr(agn + ".modules[1].moduleType", "stateMachine", type="string")
    cmds.setAttr(agn + ".modules[1].moduleName", "stateMachine", type="string")
    cmds.setAttr(agn + ".atoms_stateMachine_state", 1)

    cmds.playbackOptions(min=1, max=50)

    s = cmds.spaceLocator()[0]
    cmds.setAttr(s + ".t", 200, 50, -120)
    cmds.setKeyframe(s)
    cmds.currentTime(50)
    cmds.setAttr(s + ".t", 420, 50, 100)
    cmds.setKeyframe(s)
    cmds.currentTime(0)

    cmds.setAttr(agn + ".modules[2].moduleType", "myJointAim", type="string")
    cmds.setAttr(agn + ".modules[2].moduleName", "myJointAim", type="string")
    cmds.connectAttr(s + ".t", agn + ".atoms_myJointAim_target")
    cmds.setAttr(agn + ".atoms_myJointAim_axis", 2)
    cmds.setAttr(agn + ".atoms_myJointAim_joint", "Head", type="string")