Marching Queue

from maya import cmds

import AtomsMath
import AtomsCore
import Atoms
import AtomsUtils
from AtomsMaya.utils import agentgroup
from AtomsMaya.hostbridge.atomsgroup import MayaAtomsGroupHostBridge

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.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("pointsLayout")
    ag.add_module("stateMachine")
    ag.add_module("marchingQueue")
    
    int_type_str = AtomsCore.IntMetadata.staticTypeStr()
    str_type_str = AtomsCore.StringMetadata.staticTypeStr()
    vec_type_str = AtomsCore.Vector3Metadata.staticTypeStr()

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

    count = 6
    for i in range(count):
        ag.set_metadata_value("pointsLayout", "agentTypes", str_type_str, "man", index=i)
        ag.set_metadata_value("pointsLayout", "directions", vec_type_str, (1, 0, 0), index=i)
        ag.set_metadata_value("pointsLayout", "points", vec_type_str, ((count - i) * 60, 0, 0), index=i)
    
setup()