Editable Initial State - Attempt 13

README.md

⬅️ Back

Putting it all together, I'll borrow from the Maya devkit examples testNsolverNode and testNobjectNode for an event loop

Architecture

Roughly 800 lines in total, including comments.

• main.cpp Single translation unit, single output plug-in
• rigid.hpp The rigid, separated into its own file for readability, but is #included into main.cpp
• solver.hpp Also included, contains the solver part.

Here's the intended evaluation from start to finish, except for Parallel which doesn't play by the rules.

1. pCube is about to be rendered
2. pCube1.translateX is evaluated
3. rigid.outputTranslateX is pulled
4. solver.nextState is pulled
5. rigid.startState is pulled
6. solver steps simulation, updates the position of all rigids
7. rigid.outputTranslateX is clean
8. pCube1.translateX is clean
9. Finish

On subsequent frames, rigid.currentState is pulled instead, so we can handle initialisation once and only once.

Result

DG mode almost works, with the exception of parenting to each other. It isn't recording values from MDagMessage::addWorldMatrixModifiedCallback properly; it shouldn't happen on time change, only on manipulation at the start time.

Build & Run

From Powershell.

cd c:\
git clone https://gist.github.com/4ef2c1a8dd86d8acc61f1888afdae2d8.git initialStateNode13
mkdir initialStateNode13/build
cd initialStateNode10/build
$env:DEVKIT_LOCATION="c:\path\to\your\devkit"
cmake .. -G Ninja
ninja

Then from the Script Editor.

import os
from maya import cmds

fname = r"c:\initialState13\build\initialStateNode.mll"
cmds.loadPlugin(fname)


# Set to False for a non-cycling, functioning solver loop
USE_CONNECTION = False

# Set to True to test hierarchy
USE_PARENTING = False

time = "time1"
solver = cmds.createNode("solverNode", parent=cmds.createNode("transform", name="solver"))
cmds.setAttr(solver + ".gravity", 0, -9.81, 0)
cmds.setAttr(solver + ".startTime", cmds.getAttr(time + ".outTime"))
cmds.connectAttr(time + ".outTime", solver + ".currentTime")

parent = None
for i in range(3):
    cube, _ = cmds.polyCube()

    if parent:
        cmds.parent(cube, parent, relative=True)
        cmds.move(3, 0, 0, relative=True)
    else:
        cmds.move(-3 * i, 15, -7)

    rigid = cmds.createNode("rigidNode", parent=cube)

    cmds.connectAttr(time + ".outTime", rigid + ".currentTime")
    cmds.connectAttr(cube + ".parentInverseMatrix", rigid + ".inParentInverseMatrix")

    if USE_CONNECTION:
        cmds.connectAttr(cube + ".matrix", rigid + ".inMatrix")
    else:
        cmds.setAttr(rigid + ".inMatrix", cmds.getAttr(cube + ".matrix"), type="matrix")

    cmds.connectAttr(rigid + ".outputTranslateX", cube + ".translateX")
    cmds.connectAttr(rigid + ".outputTranslateY", cube + ".translateY")
    cmds.connectAttr(rigid + ".outputTranslateZ", cube + ".translateZ")

    # Connect to scene
    cmds.connectAttr(rigid + ".startState", solver + ".startState[%d]" % i)
    cmds.connectAttr(rigid + ".currentState", solver + ".currentState[%d]" % i)
    cmds.connectAttr(solver + ".nextState[%d]" % i, rigid + ".nextState")

    if USE_PARENTING:
        parent = cube


#cmds.file(new=True, force=True)
#cmds.unloadPlugin(os.path.basename(fname))

CMakeLists.txt

cmake_minimum_required(VERSION 2.8)

include($ENV{DEVKIT_LOCATION}/cmake/pluginEntry.cmake)

set(PROJECT_NAME initialStateNode)

set(SOURCE_FILES
   main.cpp
)

set(LIBRARIES
    OpenMaya
    OpenMayaRender
    OpenMayaUI
    Foundation
)

add_compile_definitions("/D_DEBUG /Zi")

build_plugin()

main.cpp

#include <string.h>
#include <unordered_map>
#include <cassert>

#include <maya/MPxLocatorNode.h>

#include <maya/MFnNumericAttribute.h>
#include <maya/MFnMatrixAttribute.h>
#include <maya/MFnUnitAttribute.h>
#include <maya/MFnPlugin.h>
#include <maya/MFnDagNode.h>
#include <maya/MFnTransform.h>
#include <maya/MDagPath.h>
#include <maya/MVector.h>
#include <maya/MQuaternion.h>

#include <maya/MTypeId.h> 
#include <maya/MPlug.h>
#include <maya/MString.h>
#include <maya/MVector.h>
#include <maya/MTransformationMatrix.h>
#include <maya/MMatrix.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MStreamUtils.h>
#include <maya/MNodeMessage.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnMatrixData.h>
#include <maya/MCallbackIdArray.h>
#include <maya/MGlobal.h>
#include <maya/MMessage.h>
#include <maya/MDGMessage.h>
#include <maya/MDagMessage.h>

#define Print                           MStreamUtils::stdErrorStream
#define Translate(mat)                  MTransformationMatrix(mat).translation(MSpace::kTransform)
#define AddAttribute(attr)              assert(addAttribute(attr) == MS::kSuccess);
#define AttributeAffects(attr, affects) assert(attributeAffects(attr, affects) == MS::kSuccess);

static bool debug() {
    static bool debugInfoExists { false };
    const bool debugInfo = MGlobal::optionVarIntValue("debugInfo", &debugInfoExists);
    return debugInfoExists ? debugInfo : false;
}

// Global counter, for a unique UID per rigid
static int gUniqueIdCounter { 1000 };

    
struct RigidBody {
    int uid { -1 };
    MString name;
    MVector velocity { 0, 0, 0 };
    MMatrix parentInverseMatrix {};
    MMatrix restMatrix {};
    MMatrix worldMatrix {};
    bool    worldMatrixChanged { true };

    // Callback state
    bool editable { false };
    bool edited { false };
    MMatrix editedMatrix {};
};

using RigidMap = std::unordered_map<int, RigidBody>;

// For purposes of demonstration, rather
// than passing data between nodes, we'll
// use a global variable and pass an index
// into this map between nodes instead.
static RigidMap gRigids;
static bool has(int uid) { return gRigids.find(uid) != gRigids.end(); }


// Helper functions

inline MDataHandle InputValue(MDataBlock& datablock, MObject attr) {
    MStatus status;
    MDataHandle handle = datablock.inputValue(attr, &status);
    assert(status == MS::kSuccess);
    return handle;
}

inline MDataHandle OutputValue(MDataBlock& datablock, MObject attr) {
    MStatus status;
    MDataHandle handle = datablock.outputValue(attr, &status);
    assert(status == MS::kSuccess);
    return handle;
}


#include "solver.hpp"
#include "rigid.hpp"


MStatus initializePlugin(MObject obj) { 
    MStatus   status;
    MFnPlugin plugin(obj, PLUGIN_COMPANY, "3.0", "Any");

    status = plugin.registerNode(SolverNode::name,
                                 SolverNode::id,
                                 SolverNode::creator,
                                 SolverNode::initialize,
                                 MPxNode::kLocatorNode,
                                 &SolverNode::drawDbClassification);

    status = plugin.registerNode(RigidNode::name,
                                 RigidNode::id,
                                 RigidNode::creator,
                                 RigidNode::initialize,
                                 MPxNode::kLocatorNode,
                                 &RigidNode::drawDbClassification);

    Print() << "==============\n";
    Print() << "Initialising\n";

    return status;
}


MStatus uninitializePlugin(MObject obj) {
    MStatus   status;
    MFnPlugin plugin(obj);

    status = plugin.deregisterNode(RigidNode::id);
    status = plugin.deregisterNode(SolverNode::id);

    Print() << "Deinitialising.\n";
    Print() << "==============\n";

    return status;
}

rigid.hpp

/**
 * Individual rigid body
 *
 */
class RigidNode : public MPxLocatorNode {
public:
    RigidNode() {};
    ~RigidNode() override;

    MStatus compute(const MPlug&, MDataBlock&) override;
    MStatus computeUid(const MPlug&, MDataBlock&);
    MStatus computeStartState(const MPlug&, MDataBlock&);
    MStatus computeCurrentState(const MPlug&, MDataBlock&);
    MStatus computeOutputTranslate(const MPlug&, MDataBlock&);

    static  void* creator() { return new RigidNode(); }
    static  MStatus initialize();
    void    postConstructor();
    void    reset(MDataBlock&);
    int     uid() { return _uid; }

    MStatus setDependentsDirty(const MPlug&, MPlugArray&) override;
    bool    isPassiveOutput(const MPlug&) const override;

public:
    static  MString       name;
    static  MTypeId       id;
    static const MString  drawDbClassification;
    static const MString  drawRegistrantId;

    // Attributes
    static  MObject   aUid;
    static  MObject   aCurrentTime;
    static  MObject   aInputMatrix;
    static  MObject   aInputParentInverseMatrix;
    static  MObject   aOutputTranslate;
    static  MObject   aOutputTranslateX;
    static  MObject   aOutputTranslateY;
    static  MObject   aOutputTranslateZ;

    // <-- Pull on this to fetch latest goods from the solver
    static MObject aNextState;

    // --> Deliver initial state of this rigid to the solver
    static MObject aStartState;

    // --> Deliver animated values, like damping, to the solver
    static MObject aCurrentState;

private:
    int               _uid { -1 };
    int               _lastFrame { -1 };
    MCallbackIdArray  _callbackIds;
};


MString       RigidNode::name("rigidNode");
MTypeId       RigidNode::id(0x80012);
const MString RigidNode::drawDbClassification("drawdb/geometry/RigidNode");
const MString RigidNode::drawRegistrantId("RigidNodePlugin");
MObject       RigidNode::aUid;
MObject       RigidNode::aInputMatrix;
MObject       RigidNode::aInputParentInverseMatrix;
MObject       RigidNode::aOutputTranslate;
MObject       RigidNode::aOutputTranslateX;
MObject       RigidNode::aOutputTranslateY;
MObject       RigidNode::aOutputTranslateZ;
MObject       RigidNode::aCurrentTime;
MObject       RigidNode::aNextState;
MObject       RigidNode::aStartState;
MObject       RigidNode::aCurrentState;



RigidNode::~RigidNode() {
    // Any node created will have an entry in this map
    gRigids.erase(_uid);
}


MStatus RigidNode::initialize() {
    MStatus status;
    MFnNumericAttribute numFn;
    MFnMatrixAttribute matFn;
    MFnUnitAttribute uniFn;

    aUid = numFn.create("uid", "uid", MFnNumericData::kInt, -1, &status);
    numFn.setWritable(false);
    numFn.setStorable(false);
    numFn.setKeyable(true);

    aNextState = numFn.create("nextState", "nest", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(false);

    aStartState = numFn.create("startState", "stst", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setWritable(false);

    aCurrentState = numFn.create("currentState", "cust", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setWritable(false);

    aCurrentTime = uniFn.create("currentTime", "cuti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setKeyable(true);

    aInputMatrix = matFn.create("inMatrix", "inma");
    matFn.setStorable(false);
    matFn.setReadable(false);
    aInputParentInverseMatrix = matFn.create("inParentInverseMatrix", "ipim");
    matFn.setStorable(false);
    matFn.setReadable(false);

    aOutputTranslateX = uniFn.create("outputTranslateX", "outx", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslateY = uniFn.create("outputTranslateY", "outy", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslateZ = uniFn.create("outputTranslateZ", "outz", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslate = numFn.create("outputTranslate", "outr", aOutputTranslateX, aOutputTranslateY, aOutputTranslateZ);

    AddAttribute(aUid);
    AddAttribute(aNextState);
    AddAttribute(aStartState);
    AddAttribute(aCurrentState);
    AddAttribute(aCurrentTime);
    AddAttribute(aInputMatrix);
    AddAttribute(aInputParentInverseMatrix);
    AddAttribute(aOutputTranslate);

    // Affects initial state
    AttributeAffects(aInputMatrix, aStartState);
    AttributeAffects(aInputParentInverseMatrix, aOutputTranslate);

    // Affects simulation
    AttributeAffects(aCurrentTime, aStartState);
    AttributeAffects(aCurrentTime, aCurrentState);
    AttributeAffects(aCurrentTime, aOutputTranslate);
    AttributeAffects(aNextState, aOutputTranslate);

    return status;
}


bool RigidNode::isPassiveOutput(const MPlug& plug) const {
    /**
     * Let anything connected to these still be editable
     *
     */
    return plug == aOutputTranslateX ||
           plug == aOutputTranslateY ||
           plug == aOutputTranslateZ;
}


MStatus RigidNode::setDependentsDirty(const MPlug& plugBeingDirtied, MPlugArray& affectedPlugs) {
    return MPxLocatorNode::setDependentsDirty(plugBeingDirtied, affectedPlugs);
}


static void onWorldMatrixModified(MObject& transformNode,
                                  MDagMessage::MatrixModifiedFlags& modified,
                                  void* clientData) {

    // `transformNode` isn't necessarily our plug-in, it would most
    // likely be a parent since our plug-in is a shape. But, not even
    // the immediate parent, it could also be a grandparent or beyond.

    RigidNode* node = static_cast<RigidNode*>(clientData);
    RigidBody& rigid = gRigids[node->uid()];

    if (rigid.editable) {
        MStatus status;
        MFnDagNode fn { transformNode };

        // Route: inclusiveMatrix()
        // Seems like a bad idea, going outside of the black-box
        // ------------------------
        // MDagPath dagPath;
        // fn.getPath(dagPath);
        // MMatrix matrix = dagPath.inclusiveMatrix(&status);
        // assert(status == MS::kSuccess);

        // MFnMatrixData matrixData;
        // MObject matrixObj = matrixData.create(matrix, &status);
        // assert(status == MS::kSuccess);

        // MPlug matrixPlug { node->thisMObject(), RigidNode::aInputMatrix };
        // matrixPlug.setMObject(matrixObj);

        // Route: MFnTransform
        // Bad status, unclear as to why
        // -------------------
        // MFnTransform fnt { transformNode };
        // MVector translate = fnt.getTranslation(MSpace::kWorld, &status);
        // assert(status == MS::kSuccess);
        // MQuaternion rotate;
        // status = fnt.getRotation(rotate, MSpace::kWorld);
        // assert(status == MS::kSuccess);

        // Route: worldMatrix
        // -------------------
        MPlug worldMatrixCompound { node->thisMObject(), MPxLocatorNode::worldMatrix };
        MPlug worldMatrix = worldMatrixCompound.elementByLogicalIndex(0, &status);
        assert(status == MS::kSuccess);
        MObject obj = worldMatrix.asMObject(&status);
        assert(status == MS::kSuccess);
        MFnMatrixData data(obj, &status);
        assert(status == MS::kSuccess);
        MMatrix matrix = data.matrix(&status);
        assert(status == MS::kSuccess);

        rigid.edited = true;
        rigid.editedMatrix = matrix;

        Print() << "+-- onWorldMatrixModified(): Recorded rigid[" << rigid.uid << "].translate" << Translate(matrix) << "\n";
    }
}


static void onAttributeChanged(MNodeMessage::AttributeMessage msg, MPlug& plug, MPlug& otherPlug, void* clientData) {
    if (plug == RigidNode::worldMatrix) {
        if (msg & MNodeMessage:: kAttributeEval) {
            Print() << "--- aInputMatrix evaluaeted\n";
        }
    }
}

/**
 * This isn't a good idea, we shouldn't evaluate during dirty propagation
 *
 */
// static void onPlugDirty(MObject& obj, MPlug& plug, void* clientData) {
//     if (plug == RigidNode::aInputMatrix) {
//         MStatus status { MS::kSuccess };
//         Print() << "--- onPlugDirty()\n";

//         RigidNode* node = static_cast<RigidNode*>(clientData);
//         RigidBody& rigid = gRigids.at(node->uid());

//         if (rigid.editable) {
//             Print() << "  +++ rigid.editable\n";
//             MPlug plug { obj, RigidNode::aInputMatrix };

//             MObject obj = plug.asMObject(&status);
//             assert(status == MS::kSuccess);
//             MFnMatrixData data(obj, &status);
//             assert(status == MS::kSuccess);
//             MMatrix matrix = data.matrix(&status);
//             assert(status == MS::kSuccess);
//             rigid.editedMatrix = matrix;
//             rigid.edited = true;
//         }
//     }
// }

void RigidNode::postConstructor() {
    MStatus status { MS::kSuccess };

    // Find ourselves amongst rigids hosted by the solver
    _uid = ++gUniqueIdCounter;

    // Initialise our data, in a place accessible to our solver
    gRigids[_uid] = { _uid };

    MFnDagNode fn { thisMObject() };
    MDagPath dagPath;
    fn.getPath(dagPath);

    _callbackIds.append(
        MDagMessage::addWorldMatrixModifiedCallback(
            dagPath, onWorldMatrixModified, this, &status
        )
    );

    // Bypass warning; below callback takes a reference,
    // even though it doesn't care for actually keeping
    // it around. Bad design.
    // MObject tempobj { this->thisMObject() };

    // _callbackIds.append(
    //     MNodeMessage::addNodeDirtyPlugCallback(
    //         tempobj, onPlugDirty, this, &status
    //     )
    // );
    // assert(status == MS::kSuccess);

    /**
     * This doesn't work, it won't trigger on changes to connected
     * attributes, like the pCube1.matrix
     *
     */
    // _callbackIds.append(
    //     MNodeMessage::addAttributeChangedCallback(
    //         tempobj, onAttributeChanged, this, &status
    //     )
    // );
    assert(status == MS::kSuccess);
}


void RigidNode::reset(MDataBlock& datablock) {
    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);

    MFnDagNode fn { thisMObject() };

    rigid.name = fn.name();

    if (rigid.edited) {
        Print() << "<-- Storing recorded value..\n";
        MFnMatrixData matrixData;
        MObject matrixObj = matrixData.create(rigid.editedMatrix, &status);
        assert(status == MS::kSuccess);

        MPlug matrixPlug { thisMObject(), aInputMatrix };
        status = matrixPlug.setMObject(matrixObj);
        assert(status == MS::kSuccess);

        rigid.edited = false;
        rigid.restMatrix = rigid.editedMatrix;
    }
    else {
        Print() << "--> Restoring recorded value..\n";
        MMatrix matrix = InputValue(datablock, aInputMatrix).asMatrix();
        rigid.parentInverseMatrix = InputValue(datablock, aInputParentInverseMatrix).asMatrix();
        rigid.restMatrix = matrix * rigid.parentInverseMatrix.inverse();
    }

    rigid.worldMatrix = rigid.restMatrix;

    Print() << "(---) RigidNode: Rigid[" << _uid << "] reset @ " << Translate(rigid.restMatrix) << "\n";
}


/**
 * Restore rest matrix
 *
 */
MStatus RigidNode::computeStartState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].computeStartState() ===============\n";
    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);

    // Establish position from creation
    if (_lastFrame == -1) {
        reset(datablock);
        // rigid. = true;
    }

    rigid.worldMatrix = rigid.restMatrix;
    rigid.velocity = {0, 0, 0};
    rigid.editable = true;

    MFnDagNode fn { thisMObject() };
    rigid.name = fn.name();

    OutputValue(datablock, aStartState).set(_uid);

    datablock.setClean(plug);

    return status;
}


/**
 * Read values that change over time
 *
 * We don't really have any in this example, but for example linearDamping
 *
 */
MStatus RigidNode::computeCurrentState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].computeCurrentState() ===============\n";
    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);
    rigid.editable = false;

    // Pick up the possibly-modified initial position
    if (_lastFrame == 0) {
        reset(datablock);
    }

    OutputValue(datablock, aCurrentState).set(_uid);

    datablock.setClean(plug);

    return status;
}


MStatus RigidNode::computeOutputTranslate(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode::computeOutputTranslate() -------------------------------\n";

    MStatus status { MS::kSuccess };
    RigidBody& rigid = gRigids.at(_uid);

    // Pull and let the SolverNode go to work..
    _lastFrame = InputValue(datablock, aNextState).asInt();

    Print() << "Last frame: " << _lastFrame << "\n";

    // .
    // ..
    // ...
    // aaaaand we're back.

    // At this point, the solver will have pulled on either..
    //
    // - RigidNode::computeStartState
    // - RigidNode::computeCurrentState
    //
    // ..depending on what time it was. They will in turn have populated
    // our translate values with the latest up-to-date simulated values
    // 

    // Preserve hierarchy when manipulating an object during standstill
    if (rigid.worldMatrixChanged) {
        rigid.parentInverseMatrix = InputValue(datablock, aInputParentInverseMatrix).asMatrix();
        rigid.worldMatrixChanged = false;
    }

    MMatrix localMatrix = rigid.worldMatrix * rigid.parentInverseMatrix;

    MTransformationMatrix tm { localMatrix };
    MVector translate = tm.translation(MSpace::kTransform);
    MDataHandle tx = OutputValue(datablock, aOutputTranslateX),
                ty = OutputValue(datablock, aOutputTranslateY),
                tz = OutputValue(datablock, aOutputTranslateZ);

    tx.set(translate.x);
    ty.set(translate.y);
    tz.set(translate.z);

    datablock.setClean(plug);

    Print() << "RigidNode: Rigid[" << _uid << "].translate" << translate << " written\n";

    return status;
}


MStatus RigidNode::computeUid(const MPlug& plug, MDataBlock& datablock) {
    if (debug()) Print() << "RigidNode::computeUid()\n";
    MStatus status { MS::kSuccess };

    // We'll use this from the solver mostly to sanity-check
    // that we're talking to the rigid we think we're talking to.
    MDataHandle outputHandle = OutputValue(datablock, aUid);
    outputHandle.setInt(_uid);

    datablock.setClean(plug);

    return status;
}


MStatus RigidNode::compute(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };

    // For some reason, we need to pull on time, even though we don't need it
    // in order for its dependents to continue to be dirtied when it changes.
    // The alternative is having it drawn in the channel box or AE.
    InputValue(datablock, aCurrentTime).asTime();

    if      (plug == aStartState) {
        status = computeStartState(plug, datablock);
    }

    else if (plug == aCurrentState) {
        status = computeCurrentState(plug, datablock);
    }

    else if (plug == aOutputTranslate ||
             plug == aOutputTranslateX ||
             plug == aOutputTranslateY ||
             plug == aOutputTranslateZ) {
        status = computeOutputTranslate(plug, datablock);
    }

    else if (plug == aUid) {
        status = computeUid(plug, datablock);
    }

    return status;
}

solver.hpp

/**
 * Pull on all rigid at once, and solve them in one big batch.
 * 
 */

class SolverNode : public MPxLocatorNode {
public:
    SolverNode() {};
    ~SolverNode() override;

    static  void* creator() { return new SolverNode(); }
    static  MStatus initialize();
    MStatus compute(const MPlug& plug, MDataBlock& data) override;
    MStatus computeOutputRigids(const MPlug& plug, MDataBlock& data);
    MStatus addRigids(const MPlug& plug, MDataBlock& datablock);
    MStatus updateRigids(const MPlug& plug, MDataBlock& datablock);
    MStatus setDependentsDirty(const MPlug&, MPlugArray&) override;
    void    postConstructor() override;

    MStatus stepSimulation(MDataBlock& datablock);

public:
    static MString        name;
    static MTypeId        id;
    static const MString  drawDbClassification;
    static const MString  drawRegistrantId;

    // Attributes
    static MObject   aStartTime;
    static MObject   aCurrentTime;

    static MObject   aGravity;
    static MObject   aGravityX;
    static MObject   aGravityY;
    static MObject   aGravityZ;

    // Initial state of each rigid, e.g. position and orientation
    static MObject   aStartState;

    // Current values of rigid plugs, e.g. damping
    static MObject   aCurrentState;

    // Resulting transforms after simulation
    static MObject   aNextState;

private:
    bool    _isStartTime { false };
    bool    _beforeStartTime { false };
    bool    _isFirstTime { false };
    MTime   _lastTime { };
    MTime   _currentTime { };
    double  _deltaTime { 0.0 };

    // 0-based time
    int     _startFrame { };
    int     _currentFrame { };
    int     _simulationFrame { 0 };
    MCallbackIdArray  _callbackIds;

};


MString       SolverNode::name("solverNode");
MTypeId       SolverNode::id(0x80013);
const MString SolverNode::drawDbClassification("drawdb/geometry/SolverNode");
const MString SolverNode::drawRegistrantId("SolverNodePlugin");
MObject       SolverNode::aCurrentState;
MObject       SolverNode::aStartState;
MObject       SolverNode::aNextState;
MObject       SolverNode::aStartTime;
MObject       SolverNode::aCurrentTime;
MObject       SolverNode::aGravity;
MObject       SolverNode::aGravityX;
MObject       SolverNode::aGravityY;
MObject       SolverNode::aGravityZ;


/**
 * Is this called ahead of any and all evaluation?
 *
 */
static void onForceUpdate(MTime& time, void* clientData) {
    SolverNode* node = static_cast<SolverNode*>(clientData);
    Print() << "-------------------------------------------------------------------\n";
    Print() << "(!) onForceUpdate() " << time.value() << "\n";
    Print() << "-------------------------------------------------------------------\n";
}


SolverNode::~SolverNode() {
    for (unsigned int ii=0; ii<_callbackIds.length(); ii++) 
        MMessage::removeCallback(_callbackIds[ii]);

    _callbackIds.clear();
}

MStatus SolverNode::initialize() {
    MStatus status;
    MFnNumericAttribute numFn;
    MFnUnitAttribute uniFn;

    aStartTime = uniFn.create("startTime", "stti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setKeyable(true);

    aCurrentTime = uniFn.create("currentTime", "cuti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setKeyable(true);

    aStartState = numFn.create("startState", "stst", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(true);
    numFn.setWritable(true);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aCurrentState = numFn.create("currentState", "cust", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(true);
    numFn.setWritable(true);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aNextState = numFn.create("nextState", "nest", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setWritable(false);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aGravityX = uniFn.create("gravityX", "grvx", MFnUnitAttribute::kDistance, 0.0, &status);
    aGravityY = uniFn.create("gravityY", "grvy", MFnUnitAttribute::kDistance, -981.0, &status);
    aGravityZ = uniFn.create("gravityZ", "grvz", MFnUnitAttribute::kDistance, 0.0, &status);
    aGravity = numFn.create("gravity", "grav", aGravityX, aGravityY, aGravityZ);
    numFn.setKeyable(true);

    AddAttribute(aStartTime);
    AddAttribute(aCurrentTime);
    AddAttribute(aGravity);
    AddAttribute(aCurrentState);
    AddAttribute(aStartState);
    AddAttribute(aNextState);

    AttributeAffects(aStartTime, aNextState);
    AttributeAffects(aCurrentTime, aNextState);
    AttributeAffects(aGravity, aNextState);
    AttributeAffects(aStartState, aNextState);
    AttributeAffects(aCurrentState, aNextState);

    return status;
}

void SolverNode::postConstructor() {
    MStatus status;
    _callbackIds.append(MDGMessage::addForceUpdateCallback(onForceUpdate, this, &status));
    assert(status == MS::kSuccess);
}


MStatus SolverNode::setDependentsDirty(const MPlug& plugBeingDirtied, MPlugArray& affectedPlugs) {
    return MPxLocatorNode::setDependentsDirty(plugBeingDirtied, affectedPlugs);
}


MStatus SolverNode::compute(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };

    if (plug == aNextState && plug.isElement()){
        const MTime currentTime = InputValue(datablock, aCurrentTime).asTime();
        const MTime startTime = InputValue(datablock, aStartTime).asTime();

        _currentTime = currentTime;
        _startFrame = static_cast<int>(startTime.value());
        _currentFrame = static_cast<int>(currentTime.value()) - _startFrame;
        _deltaTime = (_currentTime - _lastTime).as(MTime::kSeconds);
        _isFirstTime = _isStartTime && _deltaTime > 0.0;
        _isStartTime = currentTime == startTime;
        _beforeStartTime = currentTime < startTime;

        status = computeOutputRigids(plug, datablock);

        _lastTime = currentTime;
    }

    return status;
}


MStatus SolverNode::computeOutputRigids(const MPlug& plug, MDataBlock& datablock) {
    if (debug()) Print() << "SolverNode::computeOutputRigids()\n";
    MStatus status { MS::kSuccess };

    auto noChange = [&]() {
        MStatus status;

        const int index = plug.logicalIndex(&status);
        assert(status == MS::kSuccess);

        MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
        assert(status == MS::kSuccess);

        nextStateHandle.jumpToElement(index);
        MDataHandle outputHandle = nextStateHandle.outputValue(&status);
        assert(status == MS::kSuccess);

        outputHandle.setInt(_currentFrame);
    };

    if (_beforeStartTime) {
        noChange();
    }

    else if (_isStartTime) {
        status = addRigids(plug, datablock);
        _simulationFrame = 0;

    }

    else if (_isFirstTime) {
        status = updateRigids(plug, datablock);
    }

    else if (_deltaTime > 0.0) {
        status = updateRigids(plug, datablock);
    }

    else {
        noChange();
    }

    assert(status == MS::kSuccess);

    datablock.setClean(plug);

    return status;
}


MStatus SolverNode::addRigids(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };
    if (debug()) Print() << "SolverNode::addRigids()\n";

    // In DG mode, evaluation never happens on frames previously
    // rendered until this is called. Madness. Why oh why?
    if (debug()) Print() << "ogs -reset\n";
    MGlobal::executeCommand("ogs -reset");

    // Pull on *all* startStates
    MArrayDataHandle startStateHandle = datablock.inputArrayValue(aStartState, &status);
    assert(status == MS::kSuccess);

    MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
    assert(status == MS::kSuccess);

    const int index = plug.logicalIndex(&status);
    assert(status == MS::kSuccess);

    // Only consider rigids that are still connected
    MPlug startStatePlug { thisMObject(), aStartState };
    MPlug elementPlug = startStatePlug.elementByLogicalIndex(index, &status);
    assert(status == MS::kSuccess);

    if (!elementPlug.isConnected()) {
        Print() << "SolverNode: Disconnected index: " << index << "\n";
        return MS::kFailure; 
    }

    status = startStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);
    status = nextStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);

    // Pull on rigid.outputStartState
    int uid = startStateHandle.inputValue(&status).asInt();
    assert(status == MS::kSuccess && uid >= 0);

    Print() << "SolverNode: Solver pulled on solver.index[" << index << "] --> rigid.next[" << uid << "]\n";

    // Changed
    MDataHandle outputHandle = nextStateHandle.outputValue(&status);
    assert(status == MS::kSuccess);
    outputHandle.setInt(_currentFrame);

    return status;
}


MStatus SolverNode::updateRigids(const MPlug& plug, MDataBlock& datablock) {
    if (debug()) Print() << "SolverNode::updateRigids()\n";
    MStatus status { MS::kSuccess };

    const int index = plug.logicalIndex(&status);
    assert(status == MS::kSuccess);

    // Pull on *all* currentStates
    MArrayDataHandle currentStateHandle = datablock.inputArrayValue(aCurrentState, &status);
    assert(status == MS::kSuccess);

    MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
    assert(status == MS::kSuccess);

    status = currentStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);
    status = nextStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);

    MDataHandle inputHandle = currentStateHandle.inputValue(&status);
    assert(status == MS::kSuccess);
    inputHandle.asInt();

    if (_currentFrame > _simulationFrame) {
        stepSimulation(datablock);
    }

    // Output a "solver-changed" message (1) to Rigid
    MDataHandle outputHandle = nextStateHandle.outputValue(&status);
    assert(status == MS::kSuccess);
    outputHandle.setInt(_currentFrame);

    return status;
}


/**
 * Apply gravity and update position
 *
 */
MStatus SolverNode::stepSimulation(MDataBlock& datablock) {
    if (debug()) Print() << "SolverNode::stepSimulation()\n";
    MStatus status { MS::kSuccess };
    MVector gravity = InputValue(datablock, aGravity).asDouble3();

    for (auto& pair : gRigids) {
        RigidBody& rigid = pair.second;
        MTransformationMatrix tm { rigid.worldMatrix };

        rigid.velocity += gravity * _deltaTime;
        tm.addTranslation(rigid.velocity, MSpace::kTransform);
        Print() << "SolverNode: rigid[" << rigid.uid << "].translate " << tm.translation(MSpace::kTransform) << "\n";

        rigid.worldMatrix = tm.asMatrix();
        rigid.worldMatrixChanged = true;
    }

    _simulationFrame = _currentFrame;

    return status;
}