Marching Queue

from maya import cmds

import AtomsMath
import AtomsCore
import Atoms
import AtomsUtils


class MarchingQueueModule(Atoms.BehaviourModule):
    """
    This module assumes all agent are atomsRobots
    """

    def __init__(self):
        Atoms.BehaviourModule.__init__(self)

    def endFrame(self, agents, agroup):

        for i in range(1, len(agents)):
            skeleton = agents[i].agentType().skeleton()
            root_id = skeleton.jointId("LeftArm")
            mid_id = skeleton.jointId("LeftForeArm")
            ik_id = skeleton.jointId("LeftHand")

            if -1 in [root_id, mid_id, ik_id]:
                AtomsUtils.Logger.warning("couldn't find one or more joint ids")
                continue

            poser = AtomsCore.Poser(skeleton)
            other_pose = agents[i - 1].pose()
            wm = poser.getWorldMatrix(other_pose, root_id)
            target_pos = AtomsMath.V3d(wm[3][0], wm[3][1], wm[3][2])
            #adding a small offset to have less mesh intersection
            target_pos.y += 10

            pose = agents[i].pose()
            AtomsCore.solveTwoJointsIKNoPoleVector(pose, skeleton, target_pos,
                                                   root_id, mid_id, ik_id, 0,
                                                   True, 0.3)


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


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

    register()

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

    count = 6
    for i in range(count):
        cmds.setAttr("%s.atoms_pointsLayout_agentTypes[%d]" % (agn, i),
                     "atomsRobot", type="string")
        cmds.setAttr("%s.atoms_pointsLayout_directions[%d]" % (agn, i),
                     1, 0, 0)
        cmds.setAttr("%s.atoms_pointsLayout_points[%d]" % (agn, i),
                     (count - i) * 60, 0, 0)

    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.setAttr(agn + ".modules[2].moduleType", "marchingQueue", type="string")
    cmds.setAttr(agn + ".modules[2].moduleName", "marchingQueue", type="string")