/*
  An example of writing to an array plug
 
  Description
  -----------------------------------------------------------

  There are two ways of achieving this
 
   1. Based on input, dynamically generate plugs in an array
   2. Based on a connected array output, write to that
 
  This is an example of (2) the latter
 
 
  Usage:
  -----------------------------------------------------------
    from maya import cmds
    array = cmds.createNode("myArray")

    for i in range(5):
        emitter = cmds.createNode("transform", name="node%i" % i)
        cmds.connectAttr(array + ".values[%d]" % i, emitter + ".ty")
    
    assert cmds.getAttr(emitter + ".ty") == 4.0
*/

#include <maya/MPxNode.h>
#include <maya/MTypeId.h>
#include <maya/MString.h>
#include <maya/MPlug.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MArrayDataBuilder.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnPlugin.h>
#include <maya/MFnPluginData.h>
#include "stdio.h"  // cerr


class ArrayPlug : public MPxNode {
public:
    ArrayPlug();

    MStatus compute(const MPlug &plug, MDataBlock &dataBlock) override;
    MStatus computeValues(const MPlug &plug, MDataBlock &dataBlock);

    static void* creator() { return new ArrayPlug; }

    static MStatus initialize();
    static const MTypeId id;
    static const MString typeName;

    // Attributes
    static MObject values;
};


const MTypeId ArrayPlug::id(0x85003);
const MString ArrayPlug::typeName("myArray");

MObject ArrayPlug::values;

inline void SOFTCHECK(MStatus status, MString msg) {
    if (!status) { cerr << "ERROR: " << msg << "\n"; }
}

#define HARDCHECK(STAT, MSG)              \
    if (MS::kSuccess != STAT) {           \
        cerr << "ERROR: " << MSG << endl; \
            return MS::kFailure;          \
    }

MStatus ArrayPlug::initialize() {
    MStatus status;

    MFnTypedAttribute   typFn;
    MFnNumericAttribute numFn;

    // Each output will have a value corresponding to its logical index
    values = numFn.create("values", "v", MFnNumericData::kFloat, 0.0f, &status);
    numFn.setArray(true);
    numFn.setStorable(false);
    numFn.setReadable(true);
    numFn.setWritable(false);

    // NOTE: This is what enables us to use MArrayDataBuilder
    //       in computeValues() below
    numFn.setUsesArrayDataBuilder(true);

    SOFTCHECK(status, "failed to create values");

    addAttribute(values);

    return MStatus::kSuccess;
}


ArrayPlug::ArrayPlug() {}


/**
 * @brief Output the index or each array plug as its value
 *
 *   _____________
 *  |             |
 *  |      values o
 *  | values[0]-| o---> 0.0
 *  | values[1]-| o---> 1.0
 *  | values[2]-| o---> 2.0
 *  | values[3]-| o
 *  |             |
 *  |             |
 *  |_____________|
 *
 *
 */
MStatus ArrayPlug::computeValues(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kSuccess };

    int index = plug.logicalIndex(&status);
    HARDCHECK(status, "Could not get logical index");

    // NOTE: Rather than calling `.outputValue`
    MArrayDataHandle arrayhandle = datablock.outputArrayValue(values, &status);
    HARDCHECK(status, "computeValues : outputArrayValue");

    // NOTE: An array data handle has a different interface than interacting
    //       with a MDataHandle directly. We're essentially fetching a
    //       handle for a given index in the array. Maya calls this "builder"
    //       presumably because you can also *create* indices in this array
    //       rather than merely writing to them
    MArrayDataBuilder builder = arrayhandle.builder(&status);
    HARDCHECK(status, "computeValues : builder");

    // NOTE: Even though we're calling *add*Element we aren't actually
    //       adding an element. You can think of this as the []-syntax
    //       of a C-array, i.e. arrayhandle[index];
    MDataHandle datahandle = builder.addElement(index, &status);
    HARDCHECK(status, "computeValues : addElement");

    // Just to highlight what the value actually is
    int value = index;

    datahandle.setFloat(value);
    datablock.setClean(plug);

    return status;
}



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

    if (plug == values) {
        return computeValues(plug, datablock);
    }

    else {
        status = MS::kUnknownParameter;
    }

    return status;
}


MStatus initializePlugin (MObject obj) {
    MStatus   status;
    MFnPlugin plugin(obj, "ArrayPlug", "2020.2", "Any");

    status = plugin.registerNode(ArrayPlug::typeName,
                                 ArrayPlug::id,
                                 ArrayPlug::creator,
                                 ArrayPlug::initialize);

    if (!status) {
        status.perror("registerNode");
        return status;
    }

    return status;
}

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

    status = plugin.deregisterNode(ArrayPlug::id);

    if (!status) {
        status.perror("deregisterNode");
        return status;
    }

    return status;
}