Writing an Atoms operator (AtomsUnreal)

#include <AtomsGraph/Ports.h>
#include <Atoms/Globals.h>
#include <Atoms/Graph/Operator.h>

class JointOffsetOperator : public Atoms::Operator
{
public:

	NODE_STANDARD_MEMBERS

	JointOffsetOperator();

	virtual ~JointOffsetOperator();

	bool	compute();

private:

	AtomsGraph::PosePort* m_inPose;
	AtomsGraph::StringPort *m_jointName;
	AtomsGraph::MatrixPort *m_offsetMatrix;
	AtomsGraph::BooleanPort *m_worldSpace;
		
	int m_jointId;
	bool m_first;
};

// this must be unique
#define JOINT_OFFSET_OPERATOR_ID 9999991 

NODE_STANDARD_MEMBERS_IMPL(JointOffsetOperator)
unsigned int JointOffsetOperator::staticTypeId() { return JOINT_OFFSET_OPERATOR_ID; }
std::string JointOffsetOperator::staticTypeStr() { return std::string("JointOffsetOperator"); }

JointOffsetOperator::JointOffsetOperator() : Operator()
{
	m_jointName= new AtomsGraph::StringPort("jointName");
	m_jointName->set("");
	addInputPort(m_jointName);

	m_offsetMatrix = new AtomsGraph::MatrixPort("offsetMatrix");
	addInputPort(m_offsetMatrix);

	m_inPose = new AtomsGraph::PosePort("inPose");
	addInputPort(m_inPose);

	m_worldSpace = new AtomsGraph::BooleanPort("worldSpace");
	m_worldSpace->set(false);
	addInputPort(m_worldSpace );

	m_jointId= -1;
	m_first = true;
}

JointConstraintOperator::~JointConstraintOperator()
{
}

bool JointConstraintOperator::compute()
{
	AtomsCore::Pose &inPose = m_inPose->getRef();
	if (inPose.numJoints() == 0)
	{
		AtomsUtils::Logger::warning() << "Empty input pose";
		return false;
	}


	AtomsCore::Pose &outPose = m_outPose->getRef();
	// Copy the input pose to the out pose port
	outPose = inPose;

	// Check if the agent and the agent type are valid
	if (!m_agent) 
	{
		AtomsUtils::Logger::error() << "Invalid agent type";
		return false;
	}

	if(!m_agent->agentType())
	{
		AtomsUtils::Logger::error() << "Invalid agent type";
		return false;
	}

	// Get the skeleton from the agent type
	const AtomsCore::Skeleton& skeleton = m_agent->agentType()->skeleton();
	
	if (m_first)
	{
		m_jointId = skeleton.jointId(m_targetJoint->getRef());
	}

	if (m_jointId == -1)
	{
		Logger::error() << "Could not find " << m_targetJoint->getRef() << "joint.";
		return false;
	}


	if (worldSpace->get())
	{
		// Compute the offset in world space
		AtomsCore::Poser poser(&skeleton);
		AtomsCore::Matrix currentMtx = poser.getWorldMatrix(pose, jointId);
		poser.setWorldMatrix(pose, m_offsetMatrix->getRef()* currentMtx, jointIdTmp);
	}
	else
	{
		// Compute the offset in local space		
		AtomsCore::JointPose& jp = outPose.jointPose(m_jointId);
		jp.setMatrix(m_offsetMatrix->getRef() * jp.matrix());
	}
	return true;
}

#include <Atoms/BehaviourModule.h>


class JointOffsetModule : public Atoms::BehaviourModule
{
public:

	JointOffsetModule ();

	virtual ~JointOffsetModule ();

	void agentsCreated(const std::vector<Atoms::Agent*>& agents, Atoms::AgentGroup* agentGroup = nullptr);

	static Atoms::BehaviourModule* creator(const std::string& parameter);
};

Atoms::BehaviourModule* JointOffsetModule::creator(const std::string& parameter)
{
	return new JointOffsetModule();
}

JointOffsetModule::JointOffsetModule() :
	Atoms::BehaviourModule()
{
	AtomsCore::StringMetadata jointName("");
	addAttribute("jointName", &jointName, false);

	AtomsCore::MatrixMetadata offsetMatrix;
	addAttribute("offsetMatrix", &offsetMatrix, true);

	AtomsCore::BoolMetadata worldSpace(false);
	addAttribute("worldSpace", &worldSpace);
}

JointOffsetModule::~JointOffsetModule()
{
}

void JointOffsetModule::agentsCreated(const std::vector<Atoms::Agent*>& agents, Atoms::AgentGroup* agentGroup)
{
	AtomsCore::MapMetadata& metadata = attributes();
	const std::string& jointName= metadata.getTypedEntry<AtomsCore::StringMetadata>("jointName")->get();
	bool worldSpace = metadata.getTypedEntry<AtomsCore::BoolMetadata>("worldSpace")->get();
	AtomsCore::Matrix& offsetMatrix = metadata.getTypedEntry<AtomsCore::MatrixMetadata>("offsetMatrix")->get();

	AtomsPtr<AtomsCore::MapMetadata> offsetMatrixOverrideMap = metadata.getTypedEntry<AtomsCore::MapMetadata>("offsetMatrix_override");

		tbb::parallel_for(tbb::blocked_range<unsigned int>(0, agents.size()),
		[&](const tbb::blocked_range<unsigned int>& r)
	{
		char groupIdChar[12];
		std::string groupIdStr;
		groupIdStr.reserve(12);

		for (unsigned int i = r.begin(); i < r.end(); i++)
		{
			Atoms::Agent* agent = agents[i];
			// Get the agent dg network
			Atoms::AgentBehaviourNetwork& network = agent->network();

			// Get the last operator fron the dg
			Atoms::Operator* endOperator = network.buildPoseNode();

			// Create the joint offset operator
			Atoms::JointOffsetOperator* jointOffset = static_cast<Atoms::JointOffsetOperator*>(network.manager().createNode(Atoms::JointOffsetOperator::staticTypeStr(), "jointOffset"));
			jointOffset->setAgent(agent);

			// Connect the last operator to the joint offset operator
			network.manager().connectAttr(endOperator->name(), "outPose", jointOffset ->name(), "inPose");

			// Set the operator parameters
			jointOffset->getInputPort<AtomsGraph::StringPort>("jointName")->set(jointName);
			jointOffset->getInputPort<AtomsGraph::BooleanPort>("worldSpace")->set(worldSpace );			

			AtomsCore::Matrix agentOffsetMatrix = offsetMatrix;
			
			// Check if the offset matrix is overridden
			AtomsPtr<AtomsCore::IntMetadata> groupIdPtr = agent->metadata().getTypedEntry<AtomsCore::IntMetadata>(ATOMS_AGENT_GROUPID);
			if (groupIdPtr)
			{
				sprintf(groupIdChar, "%d", groupIdPtr->value());
				groupIdStr = groupIdChar;

				if (offsetMatrixOverrideMap)
				{
					AtomsPtr<AtomsCore::MatrixMetadata> offsetMatrixAgentMeta = offsetMatrixOverrideMap->getTypedEntry<AtomsCore::MatrixMetadata>(groupIdStr);
					if (offsetMatrixAgentMeta)
						offsetMatrix = offsetMatrixAgentMeta->get();
				}
						
			}
				
			jointOffset->getInputPort<AtomsGraph::MatrixPort>("offsetMatrix")->set(offsetMatrix);
			
			// Set the joint offset operator as the last operator to compute
			network.setBuildPoseNode(jointOffset);
		}
	});
}

void FAtomsDemoGameModule::StartupModule()
{
      AtomsUtils::Logger::info() << "Loading Joint offset plugin";
      // Register the node to the factory
      AtomsGraph::NodeFactory& manager = AtomsGraph::NodeFactory::instance();
      manager.registerNode(JointConstraintOperator::staticTypeStr(), &JointConstraintOperator::creator);

	Atoms::BehaviourModules& moduleManager = Atoms::BehaviourModules::instance();
	moduleManager.registerBehaviourModule("SimpleJointOffset", &JointOffsetModule::creator, Atoms::JointOffsetModule::kNative);
      return true;
  }
 
void FAtomsDemoGameModule::ShutdownModule()
{
    AtomsUtils::Logger::info() << "Unloading Joint offset plugin";
    // Deregister the node from the node factory
    AtomsGraph::NodeFactory& manager = AtomsGraph::NodeFactory::instance();
    manager.deregisterNode(JointConstraintOperator::staticTypeStr());
    Atoms::BehaviourModules& moduleManager = Atoms::BehaviourModules::instance();
    moduleManager.deregisterBehaviourModule("SimpleJointOffset");
    return true;
}