'''
*   This script was created by
*       Trevor Sommer
*           trevor@trevorsommer.com

Module containing numurous rigging utility or helper functions
'''

import maya.cmds as cmds
import maya.OpenMaya as OpenMaya

def isPoint3D(point):
    '''
    seemingly effective way to check if a valid point is passed as input
    '''
    try :
        test = OpenMaya.MPoint(point[0], point[1], point[2])
        return True
    except:
        return False

def isVector3D(vector):
    '''
    seemingly effective way to check if a valid Vector is passed as input
    '''
    try :
        test = OpenMaya.MVector(vector[0], vector[1], vector[2])
        return True
    except:
        return False

def vectDist(start, end):
    '''
    returns distance between two points
    '''
    try:
        startVect = OpenMaya.MVector(start[0], start[1], start[2])
        endVect = OpenMaya.MVector(end[0], end[1], end[2])
        distance = OpenMaya.MVector(endVect - startVect).length()
        return distance
    except:
        print 'One of the points passed to vectDist does not appear to be a valid 3D point'
        return False

def addSuffixToList(stringList, suffix):
    '''
    returns a new list by adding the suffix to each member of the passed stringList
    '''
    try:
        return [s + suffix for s in stringList]
    except:
        print 'addSuffixToList has failed do to an unforseen error'
        return False

def sphereControllerObj(conName):
    '''
    Creates nurbs sphere controller object primarily used for pole vector controls
    '''
    try:
        createdName = cmds.sphere(pivot= [0,0,0], startSweep= 0, endSweep= 0, radius= 0.5, degree= 3, useTolerance= 0, sections= 8, spans= 4, tolerance= 0.01, axis= [0,1,0], constructionHistory= False, name= conName)
        #disconnect from initial shading group
        cmds.disconnectAttr(("|" + createdName + "|" + createdName +"Shape.instObjGroups[0]"), initialShadingGroup.dagSetMembers[0])

        #set overrides on the shape node so it looks like a controller 
        cmds.setAttr((createdName + "Shape.overrideEnabled"), 1)
        cmds.setAttr((createdName + "Shape.overrideColor"), 17)
        return createdName
    except:
        print "Failed to Create sphere Controller object with name \"" +  conName + "\"")
        return False

def diamondControllerObj(conName):
    '''
    Creates diamond shaped nurbs Curve controller Object
    '''
    try:
        createdName = cmds.sphere(pivot= [0,0,0], startSweep= 0, endSweep= 0, radius= 0.5, degree= 1, useTolerance= 0, sections= 2, spans= 2, tolerance= 0.01, axis= [0,1,0], constructionHistory= False, name= conName)
        #disconnect from initial shading group
        cmds.disconnectAttr(("|" + createdName + "|" + createdName +"Shape.instObjGroups[0]"), initialShadingGroup.dagSetMembers[0])

        #set overrides on the shape node so it looks like a controller 
        cmds.setAttr((createdName + "Shape.overrideEnabled"), 1)
        cmds.setAttr((createdName + "Shape.overrideColor"), 17)
        return createdName
    except:
        return False

def boxControllerObj(conName):
    '''
    Creates cube controller from multiple nurbsCurves, returns newly created controller name
    '''
    if cmds.objExists(conName):
        cmds.warning('Creation of box controller \"' + conName + '\" has failed to execute properly do to object of that name already existing')
        return False
        
    try:
        cmds.group(empty= True, name= conName)

        normals = {0:[0,0,1], 1:[0,0,1], 2:[0,1,0], 3:[0,1,0]}
        centers = {0:[0,0,1], 1:[0,0,-1], 2:[0,1,0], 3:[0,-1,0]}

        for idx in range(4):
            squareName = cmds.nurbsSquare(normal= normals.get(idx), center= centers.get(idx), degree= 1, constructionHistory= 0, sideLength1= 2, sideLength2= 2, name= ("box" + str(idx)))

            for prefix in ["right", "bottom", "top", "left"]:
                cmds.parent((prefix + squareName[0] + "Shape"), conName, shape= True, relative= True,)
                cmds.rename((prefix+ squareName[0] + "Shape"),(conName + prefix + squareName[0] + "Shape"))

            cmds.delete(squareName[0])
        return conName
    except:
        cmds.warning('Creation of box controller \"' + conName + '\" has failed to execute properly')
        return False
    
def getShape(_object):
    '''
    returns shape nodes of passed object if they exists
    '''
    try:
        shapes = cmds.listRelatives(_object, shapes=True)
        if shapes:
            return shapes
    except:
        cmds.warning('The object \"' + _object + '\" does not have any shape nodes')
        return False       

def polePosCalc(start, end, Magnitude=5.0):
    '''
    returns a point position, offset by magnitude, for rigging placement/ik/fk switch to match the pole vector 
    '''
    try:
        startPos = cmds.pointPosition((start + '.rotatePivot'), world=True)
        endPos = cmds.pointPosition((end + '.rotatePivot'), world=True)

        startVect = OpenMaya.MVector(startPos[0], startPos[1], startPos[2])
        endVect = OpenMaya.MVector(endPos[0], endPos[1], endPos[2])

        targetVector = OpenMaya.MVector(endVect - startVect).normal()
        finalPoint = OpenMaya.MPoint(endPos[0], endPos[1], endPos[2]) + (targetVector * Magnitude)
        outPoint = [finalPoint.x, finalPoint.y, finalPoint.z]
        return outPoint

    except:
        print 'One of the objects passed to polePosCalc does not appear to be a valid 3D point'
        return False

def attrLocker(object, attrList, lockAttr=True, keyable=False):
    if not cmds.objExists(object):
        print ('No object by the name \"' + object + '\" exists Canceling attrLocker')
        return False
    if not isinstance(lockAttr, bool) or not isinstance(keyable, bool):
        print "attrLocker requires \"lockAttr\" and \"keyable\" arguements to be of type bool, canceling operation"
        return False
    if not isinstance(attrList, str) and not all(isinstance(item, str) for item in attrList):
        print "attrLocker requires you pass it a list of strings that represent the attributes to lock or hide, canceling operation"
        return False

    for attribute in attrList:
        if cmds.attributeQuery(attribute, exists= True, node= object):
            cmds.setAttr((object + "." + attribute), lock= lockAttr, keyable= keyable)

    return True

def get_charPfx(nameList):
    '''
    Returns the charPfx if the criteria for character prefix exists on the passed string/stringlist 
    will check all strings in the list against one another and if any inconsistancy in the character prefix is found will return False
    charPfx is defined as = (char + "_" + (if any ["lf","cn","rt"]))
    '''
    print nameList
    if not isinstance(nameList, str) and not all(isinstance(item, basestring) for item in nameList):
        print "No valid Character Prefix was found"
        return False

    if not isinstance(nameList, (list, tuple)):
        nameList = [nameList]

    allCharPfx = []

    for item in nameList:
        currentCharPfx = ""
        subStrings = item.split("_")
        if len(subStrings) < 2:
            print "No Consistant valid Character Prefix was found"
            return False

        if item.startswith((subStrings[0] + "_lf")):
            currentCharPfx = subStrings[0]
        elif item.startswith((subStrings[0] + "_rt")):
            currentCharPfx = subStrings[0]
        elif item.startswith((subStrings[0] + "_cn")):
            currentCharPfx = subStrings[0]
        else:
            print subStrings[0]
            print "No Consistant valid Character Prefix was found"
            return False        

        if allCharPfx and currentCharPfx not in allCharPfx:
            print "No Consistant valid Character Prefix was found"
            return False 
        else:
            allCharPfx.append(currentCharPfx)

    return allCharPfx[0]