nodes_v2.py

'''

'''

from collections import namedtuple

import maya.cmds as cmds
import maya.OpenMaya as om
import maya.api.OpenMaya as om2


def connect(src, dest):
    
    def getPlug(attr):
        sel = om2.MSelectionList()
        sel.add(attr)
        return sel.getPlug(0)
        
    if isinstance(src, (int, float, bool)) and cmds.objExists(str(dest)):
        return cmds.setAttr(str(dest), src)
    if isinstance(dest, (list, tuple)):
        if isinstance(src, (list, tuple)):
            return [connect(src[i], dest[i]) for i in range(min(len(src), len(dest)))]
        if cmds.objExists(str(src)):
            return [connect(getPlug(src).child(i), dest[i]) for i in range(min(getPlug(src).numChildren(), len(dest)))]
    destPlug = getPlug(dest)
    if isinstance(src, (int, float, bool)) and dest.isCompound:
        return [connect(src, destPlug.child(i)) for i in range(destPlug.numChildren())]
    if isinstance(src, (list, tuple)) and destPlug.isArray:
        return [connect(src[i], destPlug.elementByLogicalIndex(i)) for i in range(min(len(value), destPlug.numElements()))]
    if isinstance(src, (list, tuple)) and destPlug.isCompound:
        return [connect(src[i], destPlug.child(i)) for i in range(min(len(src), destPlug.numChildren()))]
    if isinstance(src, (list, tuple)) and destPlug.isArray:
        return connect(src, destPlug.elementByLogicalIndex(0))
    if getPlug(src).isCompound:
        return cmds.connectAttr(str(src), str(dest), force=True)
    if destPlug.isArray:
        return connect(src, destPlug.elementByLogicalIndex(0))
    if destPlug.isCompound:
        return [connect(src, destPlug.child(i)) for i in range(destPlug.numChildren())]
    return cmds.connectAttr(str(src), str(dest), force=True)

def connectOrSet(source, destination):
    '''
    Either sets the destination to the source value or connects the
    source attr to the destination attr.
    :param source: [str or value]
    :param destination: [str]
    '''

    def listChildren(input):
        node = str(input).split('.')[0]
        attr = str(input).split('.')[-1]

        children = cmds.attributeQuery(attr.split('[')[0], node=node, listChildren=True)
        isMulti = cmds.attributeQuery(attr.split('[')[0], node=node, multi=True)
        if children is not None and not isMulti:
            return ['%s.%s' % (node, child) for child in children]
        elif children is not None:
            return ['.'.join([str(input), child]) for child in children]
        else:
            return []

    def getType(input):
        if isinstance(input, basestring) or cmds.objExists(str(input)):
            node = str(input).split('.')[0]
            attr = str(input).split('.')[-1]
            if '[' in attr:
                attr = attr.split('[')[0]
            elif cmds.attributeQuery(attr, node=node, multi=True):
                return 'arrayAttribute'

            if cmds.attributeQuery(attr, node=node, listChildren=True) is not None:
                return 'compoundAttribute'
            else:
                return 'attribute'
        elif isinstance(input, list):
            return 'list'
        elif isinstance(input, int) or isinstance(input, float) or isinstance(input, bool):
            return 'value'
        else:
            raise NotImplementedError('Input type %s is not supported' % str(type(input)))

    types = (getType(source), getType(destination))
    if types == ('value', 'attribute'):
        cmds.setAttr(str(destination), source)

    elif types == ('compoundAttribute', 'compoundAttribute') or types == ('attribute', 'attribute'):
        cmds.connectAttr(str(source), str(destination), force=True)


    elif types == ('value', 'compoundAttribute'):
        destinationChildren = listChildren(destination)
        for child in destinationChildren:
            connectOrSet(source, child)

    elif types == ('list', 'list'):
        for i in range(min(len(source), len(destination))):
            connectOrSet(source[i], destination[i])

    elif types == ('list', 'arrayAttribute') or types == ('list', 'compoundAttribute'):
        destinationChildren = listChildren(destination)
        for i in range(min(len(source), len(destinationChildren))):
            connectOrSet(source[i], destinationChildren[i])

    elif types == ('compoundAttribute', 'list'):
        sourceChildren = listChildren(source)
        for i in range(min(len(sourceChildren), len(destination))):
            connectOrSet(sourceChildren[i], destination[i])

    elif types == ('attribute', 'arrayAttribute') or types == ('value', 'arrayAttribute'):
        connectOrSet(source, destination + '[0]')

    elif types == ('attribute', 'compoundAttribute'):
        destinationChildren = listChildren(destination)
        for i in range(len(destinationChildren)):
            connectOrSet(source, destinationChildren[i])

    else:
        raise NotImplementedError('Cannot connect %s %s to %s %s' % (types[0], source, types[1], destination))


'''
Basic Nodes
'''

def add(*values, **kwargs):
    '''
    Returns the attribute sum of the input values.
    :param values: [value or attribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'Add'
    node = cmds.createNode('plusMinusAverage', name=name)
    for i, value in enumerate(values):
        connectOrSet(value, '%s.input1D[%s]' % (node, i))
    return '%s.output1D' % node


def subtract(*values, **kwargs):
    '''
    Returns the attribute difference of the input values.
    :param values: [value or attribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'Subtract'
    node = cmds.createNode('plusMinusAverage', name=name)
    cmds.setAttr('%s.operation' % node, 2)
    for i, value in enumerate(values):
        connectOrSet(value, '%s.input1D[%s]' % (node, i))
    return '%s.output1D' % node


def average(*values, **kwargs):
    '''
    Returns the attribute average of the input values.
    :param values: [value or attribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'Average'
    node = cmds.createNode('plusMinusAverage', name=name)
    cmds.setAttr('%s.operation' % node, 3)
    for i, value in enumerate(values):
        connectOrSet(value, '%s.input1D[%s]' % (node, i))
    return '%s.output1D' % node


def multiply(input1, input2, name='Multiply'):
    '''
    Multiplys the two input values.
    :param input1: [value or attribute]
    :param input2: [value or attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('multiplyDivide', name=name)
    connectOrSet(input1, '%s.input1X' % node)
    connectOrSet(input2, '%s.input2X' % node)
    return '%s.outputX' % node


def divide(input1, input2, name='Divide'):
    '''
    Divides the two input values.
    :param input1: [value or attribute]
    :param input2: [value or attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('multiplyDivide', name=name)
    cmds.setAttr('%s.operation' % node, 2)
    connectOrSet(input1, '%s.input1X' % node)
    connectOrSet(input2, '%s.input2X' % node)
    return '%s.outputX' % node


def pow(value, exponent=2.0, name='Power'):
    '''
    Raises the input value to the power of the given exponent. Default squares the input.
    :param value: [value or attribute]
    :param exponent: [value or attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('multiplyDivide', name=name)
    cmds.setAttr('%s.operation' % node, 3)
    connectOrSet(value, '%s.input1X' % node)
    connectOrSet(exponent, '%s.input2X' % node)
    return '%s.outputX' % node


def sqrt(value, name='SquareRoot'):
    '''
    Returns the attribute square root of the given value.
    :param value: [value or attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('multiplyDivide', name=name)
    cmds.setAttr('%s.operation' % node, 3)
    connectOrSet(value, '%s.input1X' % node)
    connectOrSet(0.5, '%s.input2X' % node)
    return '%s.outputX' % node


def reverse(value, name='Reverse'):
    '''
    Returns the value subtracted from one.
    :param value: [value or attribute]
    :return: [str]
    '''
    node = cmds.createNode('reverse', name=name)
    connectOrSet(value, '%s.inputX' % node)
    return '%s.outputX' % node


def clamp(value, min=0, max=1, name='Clamp'):
    '''
    Return the attribute clamped.
    :param value: [value or attribute]
    :param min: [value or attribute]
    :param max: [value or attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('clamp', name=name)
    connectOrSet(value, '%s.inputR' % node)
    connectOrSet(min, '%s.minR' % node)
    connectOrSet(max, '%s.maxR' % node)
    return '%s.outputR' % node


def maxValue(input1, input2, name='Max'):
    '''
    Returns the greater attribute.
    :param input1: [Attribute or Value]
    :param input2: [Attribute or Value]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('condition', name=name)
    cmds.setAttr('%s.operation' % node, 2)
    connectOrSet(input1, '%s.firstTerm' % node)
    connectOrSet(input1, '%s.colorIfTrueR' % node)
    connectOrSet(input2, '%s.secondTerm' % node)
    connectOrSet(input2, '%s.colorIfFalseR' % node)
    return '%s.outColorR' % node


def minValue(input1, input2, name='Min'):
    '''
    Returns the lesser attribute.
    :param input1: [Attribute or Value]
    :param input2: [Attribute or Value]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('condition', name=name)
    cmds.setAttr('%s.operation' % node, 4)
    connectOrSet(input1, '%s.firstTerm' % node)
    connectOrSet(input1, '%s.colorIfTrueR' % node)
    connectOrSet(input2, '%s.secondTerm' % node)
    connectOrSet(input2, '%s.colorIfFalseR' % node)
    return '%s.outColorR' % node


def roundValue():
    pass


def floor():
    pass


def ceil():
    pass


'''
Vector Nodes
'''

def addVector(*vectors, **kwargs):
    '''
    Returns the attribute sum of the input vectors.
    :param vectors: [list or compoundAttribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'AddVector'
    node = cmds.createNode('plusMinusAverage', name=name)
    for i, vector in enumerate(vectors):
        connectOrSet(vector, '%s.input3D[%s]' % (node, i))
    return '%s.output3D' % node


def subtractVector(*vectors, **kwargs):
    '''
    Returns the attribute difference of the input vectors.
    :param vectors: [list or compoundAttribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'SubtractVector'
    node = cmds.createNode('plusMinusAverage', name=name)
    cmds.setAttr('%s.operation' % node, 2)
    for i, vector in enumerate(vectors):
        connectOrSet(vector, '%s.input3D[%s]' % (node, i))
    return '%s.output3D' % node


def averageVector(*vectors, **kwargs):
    '''
    Returns the attribute average of the input vectors.
    :param vectors: [list or compoundAttribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'AverageVector'
    node = cmds.createNode('plusMinusAverage', name=name)
    cmds.setAttr('%s.operation' % node, 3)
    for i, vector in enumerate(vectors):
        connectOrSet(vector, '%s.input3D[%s]' % (node, i))
    return '%s.output3D' % node


def multiplyVector(input1, input2, name='MultiplyVector'):
    '''
    Multiplies the components of each input vector.
    :param input1: [list or compoundAttribute]
    :param input2: [list or compoundAttribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('multiplyDivide', name=name)
    connectOrSet(input1, '%s.input1' % node)
    connectOrSet(input2, '%s.input2' % node)
    return '%s.output' % node


def normalizeVector(vector, name='NormalizeVector'):
    '''
    Normalizes the input vector.
    :param vector: list[attribute or value] or compoundAttribute
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('vectorProduct', name=name)
    cmds.setAttr('%s.operation' % node, 0)
    cmds.setAttr('%s.normalizeOutput' % node, True)
    connectOrSet(vector, '%s.input1' % node)
    return '%s.output' % node


def dotProduct(input1, input2, name='DotProduct'):
    '''
    The dot product of the given attributes or values.
    :param input1: list[attribute or value] or compoundAttribute
    :param input2: list[attribute or value] or compoundAttribute
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('vectorProduct', name=name)
    connectOrSet(input1, '%s.input1' % node)
    connectOrSet(input2, '%s.input2' % node)
    return '%s.outputX' % node


def crossProduct(input1, input2, normalize=True, name='CrossProduct'):
    '''
    The cross product of the given attributes or values.
    :param input1: list[attribute or value] or compoundAttribute
    :param input2: list[attribute or value] or compoundAttribute
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('vectorProduct', name=name)
    cmds.setAttr('%s.operation' % node, 2)
    cmds.setAttr('%s.normalizeOutput' % node, normalize)
    connectOrSet(input1, '%s.input1' % node)
    connectOrSet(input2, '%s.input2' % node)
    return '%s.output' % node


def blendVector(vector1, vector2, blend, name='BlendVector'):
    '''
    Returns the attribute blend between the two input vector attributes or values.
    :param vector1: list[attribute or value] or compoundAttribute
    :param vector2: list[attribute or value] or compoundAttribute
    :param blend: [attribute or value]
    :param name: [str]
    :return: [str]
    '''

    node = cmds.createNode('blendColors', name=name)
    connectOrSet(vector1, '%s.color1' % node)
    connectOrSet(vector2, '%s.color2' % node)
    connectOrSet(blend, '%s.blender' % node)
    return '%s.output' % node


def distance(point1, point2, name='Distance'):
    '''
    Returns the attribute distance between two points.
    :param point1: list[attribute or value] or compoundAttribute
    :param point2: list[attribute or value] or compoundAttribute
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('distanceBetween', name=name)
    connectOrSet(point1, '%s.point1' % node)
    connectOrSet(point2, '%s.point2' % node)
    return '%s.distance' % node


'''
Matrix Nodes
'''

def fourByFourMatrix(*rows, **kwargs):
    '''
    Creates a 4x4 matrix given each row.
    :param rows: list[attribute or value] or OpenMaya.MMatrix
    :return: str
    '''
    name = kwargs['name'] if 'name' in kwargs else 'Matrix'
    node = cmds.createNode('fourByFourMatrix', name=name)

    if isinstance(rows[0], om.MMatrix):
        rows = [list(row) for row in rows[0]]
    elif isinstance(rows[0], om2.MMatrix):
        rows = [[rows[0].getElement(i, j) for j in range(4)] for i in range(4)]

    for i, row in enumerate(rows):
        connectOrSet(row, ['%s.i%s%s' % (node, i, j) for j in range(4)])

    return '%s.output' % node


def decomposeMatrix(inputMatrix, name='DecomposeMatrix'):
    '''
    Decomposes a given matrix attribute.
    :param inputMatrix: [str]
    :param name: [str]
    :return: [namedtuple] A collection of the separate transformation values.
    '''
    node = cmds.createNode('decomposeMatrix', name=name)
    connectOrSet(inputMatrix, '%s.inputMatrix' % node)
    DecomposeData = namedtuple('DecomposeData', ['translate', 'rotate', 'scale', 'shear', 'quaternion'])
    return DecomposeData(
        translate='%s.outputTranslate' % node, rotate='%s.outputRotate' % node, scale='%s.outputScale' % node,
        shear='%s.outputShear' % node, quaternion='%s.outputQuat' % node
    )


def multiplyMatrix(*matrices, **kwargs):
    '''
    Returns the sum of each input matrix.
    :param matrices: [attribute]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'MultiplyMatrix'
    node = cmds.createNode('multMatrix', name=name)
    for i, matrix in enumerate(matrices):
        connectOrSet(matrix, '%s.matrixIn[%s]' % (node, i))
    return '%s.matrixSum' % node


def transposeMatrix():
    pass


def blendMatrix(*matrixWeights, **kwargs):
    '''
    Returns a weighted sum of each input matrix, weight pair.
    :param matrixWeights: list[attribute, attribute or value]
    :return: [str]
    '''
    name = kwargs['name'] if 'name' in kwargs else 'BlendMatrix'
    node = cmds.createNode('wtAddMatrix', name=name)
    for i, matrixWeight in enumerate(matrixWeights):
        connectOrSet(matrixWeight[0], '%s.wtMatrix[%s].matrixIn' % (node, i))
        connectOrSet(matrixWeight[1], '%s.wtMatrix[%s].weightIn' % (node, i))
    return '%s.matrixSum' % node


def inverseMatrix(inputMatrix, name='InverseMatrix'):
    '''
    Returns the inverted input matrix attribute.
    :param inputMatrix: [attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('inverseMatrix', name=name)
    connectOrSet(inputMatrix, '%s.inputMatrix' % node)
    return '%s.outputMatrix' % node


def composeMatrix(translate=None, rotate=None, scale=None, quaternion=None, shear=None, rotateOrder=None,
                  name='ComposeMatrix'):
    '''
    Creates a matrix attribute given input transformations.
    :param translate: [str]
    :param rotate: [s]tr
    :param scale: [str]
    :param quaternion: [str]
    :param shear: [str]
    :param rotateOrder: [int]
    :return: [str]
    '''
    node = cmds.createNode('composeMatrix', name=name)
    if translate is not None:
        connectOrSet(translate, '%s.inputTranslate' % node)
    if rotate is not None:
        connectOrSet(rotate, '%s.inputRotate' % node)
    if scale is not None:
        connectOrSet(scale, '%s.inputScale' % node)
    if shear is not None:
        connectOrSet(shear, '%s.inputShear' % node)
    if quaternion is not None:
        connectOrSet(quaternion, '%s.inputQuat' % node)
        cmds.setAttr('%s.useEulerRotation' % node, False)
    if rotateOrder is not None:
        connectOrSet(rotateOrder, '%s.inputRotateOrder' % node)
    return '%s.outputMatrix' % node


def decomposeTranslate(inputMatrix, name='DecomposeTranslate'):
    '''
    :param inputMatrix: [str]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('decomposeMatrix', name=name)
    connectOrSet(inputMatrix, '%s.inputMatrix' % node)
    return '%s.outputTranslate' % node


def decomposeRotate(inputMatrix, name='DecomposeRotate'):
    '''
    :param inputMatrix: [str]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('decomposeMatrix', name=name)
    connectOrSet(inputMatrix, '%s.inputMatrix' % node)
    return '%s.outputRotate' % node


def decomposeScale(inputMatrix, name='DecomposeScale'):
    '''
    :param inputMatrix: [str]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('decomposeMatrix', name=name)
    connectOrSet(inputMatrix, '%s.inputMatrix' % node)
    return '%s.outputScale' % node


def vectorMatrixProduct(inputVector, inputMatrix, normalize=False, name='VectorMatrixProduct'):
    '''
    Multiplys a given vector by a given matrix.
    :param inputVector: list[attribute or value] or compoundAttribute
    :param inputMatrix: [attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('vectorProduct', name=name)
    cmds.setAttr('%s.operation' % node, 3)
    cmds.setAttr('%s.normalizeOutput' % node, normalize)
    connectOrSet(inputVector, '%s.input1' % node)
    connectOrSet(inputMatrix, '%s.matrix' % node)
    return '%s.output' % node


def pointMatrixProduct(inputVector, inputMatrix, normalize=False, name='PointMatrixProduct'):
    '''
    Multiplys a given point by a given matrix.
    :param inputVector: list[attribute or value] or compoundAttribute
    :param inputMatrix: [attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('vectorProduct', name=name)
    cmds.setAttr('%s.operation' % node, 4)
    cmds.setAttr('%s.normalizeOutput' % node, normalize)
    connectOrSet(inputVector, '%s.input1' % node)
    connectOrSet(inputMatrix, '%s.matrix' % node)
    return '%s.output' % node


def matrixDistance(point1, point2, name='Distance'):
    '''
    Returns the attribute distance between two points.
    :param point1: list[attribute or value] or compoundAttribute
    :param point2: list[attribute or value] or compoundAttribute
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('distanceBetween', name=name)
    connectOrSet(point1, '%s.inMatrix1' % node)
    connectOrSet(point2, '%s.inMatrix2' % node)
    return '%s.distance' % node


'''
Rotation
'''

def eulerToQuaternion(inputRotate, name='EulerToQuat'):
    node = cmds.createNode('eulerToQuat', name=name)
    connectOrSet(inputRotate, '%s.inputRotate' % node)
    return ['%s.outputQuat%s' % (node, axis) for axis in ['X', 'Y', 'Z', 'W']]


'''
Conditions
'''

def choice(selector, choices, name='Choice'):
    '''
    Returns an input attribute given a list of choices and an index.
    :param selector: [Attibute or Value]
    :param choices: list[Attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('choice', name=name)
    connectOrSet(selector, '%s.selector' % node)
    for i, choice in enumerate(choices):
        connectOrSet(choice, '%s.input[%s]' % (node, i))

    return '%s.output' % node


def isGreater():
    pass


def isLessThan():
    pass


def isEqual(firstTerm, secondTerm, name='isEqual'):
    '''
    Returns 1 if the terms are equal, otherwise zero.
    :param firstTerm: [value or Attribute]
    :param secondTerm: [value or Attribute]
    '''
    node = cmds.createNode('condition', name=name)
    cmds.setAttr('%s.operation' % node, 0)
    cmds.setAttr('%s.colorIfTrueR' % node, 1)
    cmds.setAttr('%s.colorIfFalseR' % node, 0)
    connectOrSet(firstTerm, '%s.firstTerm' % node)
    connectOrSet(secondTerm, '%s.secondTerm' % node)

    return '%s.outColorR' % node


'''
NURBS
'''

def arcLength(curve, name='ArcLength'):
    '''
    Returns the attribute arclength of the given curve.
    :param curve: [attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('curveInfo', name=name)
    connectOrSet(curve, '%s.inputCurve' % node)
    return '%s.arcLength' % node


def curveInfo(curve, name='CurveInfo'):
    '''
    Returns a container of various information about a curve.
    :param curve: [attribute]
    :param name: [str]
    :return: [str]
    '''
    node = cmds.createNode('curveInfo', name=name)
    connectOrSet(curve, '%s.inputCurve' % node)
    CurveInfo = namedtuple('CurveInfo', ['arcLength', 'controlPoints', 'knots', 'weights'])
    return CurveInfo(
        arcLength='%s.arcLength' % node, controlPoints='%s.controlPoints' % node,
        knots='%s.knots' % node, weights='%s.weights' % node
    )


def pointOnCurve(curve, parameter=0, name='PointOnCurveInfo'):
    pass


def pointOnSuface(surface, u=0, v=0, name='PointOnSurfaceInfo'):
    pass


def closestPointOnCurve(point, curve, name='ClosestPointOnCurve'):
    pass


def closestPointOnSurface(point, surface, name='ClosestPointOnSurface'):
    pass


def closestPointOnMesh(point, mesh, name='ClosestPointOnMesh'):
    pass