import pymel.core as pm
import maya.OpenMaya as OpenMaya
import traceback
import pickle
# import cPickle as pickle

# import atcore.atvc.atvc as atvc
# import atcore.atmaya.utils as utils

# # init versionControl object
# vc = atvc.VersionControl()


def getShape(node, intermediate=False):
    '''
    Gets the shape of a given node
    Returns None if unable to find shape
    Returns shape if found
    :param node:
    :param intermediate:
    :return:
    '''

    if pm.nodeType(node) == "transform":
        shapeNodes = pm.listRelatives(node, shapes=True, path=True)

        if not shapeNodes:
            shapeNodes = []

        for shapeNode in shapeNodes:
            isIntermediate = pm.getAttr("%s.intermediateObject" % shapeNode)

            if intermediate and isIntermediate and pm.listConnections(shapeNode, source=False):
                return shapeNode

            elif not intermediate and not isIntermediate:
                return shapeNode

        if shapeNodes:
            return shapeNodes[0]

    elif pm.nodeType(node) in ["mesh", "nurbsCurve", "nurbsSurface"]:
        return pm.PyNode(node)

    return None


def getShapes(nodeList):
    '''
    Runs getShape on a list of nodes
    :param nodeList:
    :return: <list>
    '''

    shapeList = []

    for node in nodeList:
        shapeNode = getShape(node)
        if shapeNode:
            shapeList.append(shapeNode)

    return shapeList


def getTransform(node):
    '''
    Gets the Transform node if a given node
    Returns None if no Transform was found
    Returns Transform object if found
    :param node:
    :return:
    '''

    node = pm.PyNode(node)
    if node.type() == 'transform':
        lhistory = node.listHistory(type='mesh')
        if len(lhistory) > 0:
            return node
    elif node.type() == "mesh":
        transNode = node.getParent()
        if transNode:
            return transNode
    return None


def getTransforms(nodeList):
    '''
    Gets the Transform nodes for all nodes in a given list
    returns list of nodes
    :param nodeList:
    :return: <list>
    '''

    transList = []

    for node in nodeList:
        transNode = getTransform(node)
        if transNode:
            transList.append(transNode)

    return transList

def getMDagPath(nodeName):
   '''
   Helper function to create an MDagPath for a given object name
   :param nodeName: <string> name of node
   :return: <MDagPath>
   '''

   selectionList = OpenMaya.MSelectionList()
   selectionList.add(str(nodeName))
   objDagPath = OpenMaya.MDagPath()
   selectionList.getDagPath(0, objDagPath)
   return objDagPath


def getBarycentricWeights(srcMesh, tgtMesh, flipX=False):
    '''
    Builds a barycentric weight dictionary for a given source and target mesh
    Iterates over each vtx in the target mesh finding the relation to the closest point and triangle on the source mesh
    Dict structure:
    {targetVtxId: {srcTriVtxId1: weight, srcTriVtxId2: weight, srcTriVtxId3: weight}}
    FlipX allows you to flip the target mesh point array over the world X-axis. Useful for mirror functions.
    :param srcMesh: shapeNode
    :param tgtMesh: shapeNode
    :param flipX: <bool>
    :return: <dictionary>
    '''
    baryWeightDict = {}

    if pm.nodeType(srcMesh) != "mesh" and pm.nodeType(tgtMesh) != "mesh":
        srcMesh = getShape(srcMesh)
        tgtMesh = getShape(tgtMesh)

    try:
        srcMeshDagPath = getMDagPath(srcMesh)
        tgtMeshDagPath = getMDagPath(tgtMesh)
    except:
        print(traceback.format_exc())
        return

    # create mesh iterator
    comp = OpenMaya.MObject()
    currentFace = OpenMaya.MItMeshPolygon(srcMeshDagPath, comp)

    # get all points from target mesh
    meshTgtMPointArray = OpenMaya.MPointArray()
    meshTgtMFnMesh = OpenMaya.MFnMesh(tgtMeshDagPath)
    meshTgtMFnMesh.getPoints(meshTgtMPointArray, OpenMaya.MSpace.kWorld)

    # create mesh intersector
    matrix = srcMeshDagPath.inclusiveMatrix()
    node = srcMeshDagPath.node()
    intersector = OpenMaya.MMeshIntersector()
    intersector.create(node, matrix)

    # create variables to store the returned data
    pointInfo = OpenMaya.MPointOnMesh()
    uUtil, vUtil = OpenMaya.MScriptUtil(0.0), OpenMaya.MScriptUtil(0.0)
    uPtr = uUtil.asFloatPtr()
    vPtr = vUtil.asFloatPtr()
    pointArray = OpenMaya.MPointArray()
    vertIdList = OpenMaya.MIntArray()

    # dummy variable needed in .setIndex()
    dummy = OpenMaya.MScriptUtil()
    dummyIntPtr = dummy.asIntPtr()

    # For each point on the target mesh
    # Find the closest triangle on the source mesh.
    # Get the vertIds and the barycentric coords.

    #
    if flipX:
        meshTgtMPointArray = flipMPointArrayInX(meshTgtMPointArray)

    for i in range(meshTgtMPointArray.length()):

        intersector.getClosestPoint(meshTgtMPointArray[i], pointInfo)
        pointInfo.getBarycentricCoords(uPtr, vPtr)
        u = uUtil.getFloat(uPtr)
        v = vUtil.getFloat(vPtr)

        faceId = pointInfo.faceIndex()
        triangleId = pointInfo.triangleIndex()

        currentFace.setIndex(faceId, dummyIntPtr)
        currentFace.getTriangle(triangleId, pointArray, vertIdList, OpenMaya.MSpace.kWorld)

        weightDict = {}
        weightDict[vertIdList[0]] = u
        weightDict[vertIdList[1]] = v
        weightDict[vertIdList[2]] = 1 - u - v

        baryWeightDict[i] = weightDict

    return baryWeightDict


def flipMPointArrayInX(MPointArray):
    '''
    Flips an MPoint array in X
    :param MPointArray: <MPointArray>
    :return:
    '''

    flippeMPointArray = OpenMaya.MPointArray()

    for i in range(MPointArray.length()):
        pX = MPointArray[i][0] * -1
        pY = MPointArray[i][1]
        pZ = MPointArray[i][2]
        flippeMPointArray.append(OpenMaya.MPoint(pX, pY, pZ))

    return flippeMPointArray


def findVertsOnSidesOfX(MPointArray):
    '''
    Bins an MPointArray into two lists of verticeIndexes.
    Points living within -0.0001 to 0.0001 on the X axis are discarded.
    :param MPointArray: <MPointArray>
    :param positiveSide: <bool>
    :return: <list> <list> positiveVtxs, negativeVtxs
    '''

    positiveVtxs = []
    negativeVtxs = []

    for i in range(MPointArray.length()):
        pX = MPointArray[i][0]

        # not equal to 0 as we are not interested in vertices in the "perfect middle"
        if pX > 0.0001:
            positiveVtxs.append(i)
        elif pX < -0.0001:
            negativeVtxs.append(i)

    return positiveVtxs, negativeVtxs


def getMPointArray(MDagPath):
    '''
    Gets an MPointArray for all points in a mesh.
    MPoint is usable like a python list in the structure
    MPoint[index][x-pos, y-pos, z-pos]
    :param MDagPath: <MDagPath>
    :return:
    '''

    MPointArray = OpenMaya.MPointArray()
    meshTgtMFnMesh = OpenMaya.MFnMesh(MDagPath)
    meshTgtMFnMesh.getPoints(MPointArray, OpenMaya.MSpace.kWorld)

    return MPointArray


def zipClusterWeights(baryWeightDict, clusterWeightDict):
    '''
    Zip a barycentric weight dictionary with a source cluster weight dictionary.
    By iterating over each target vertice in the baryWeightDict, build resulting weights, store in zippedCluster
    :param baryWeightDict: <dict> weight dict structured by the Utils modules .getBarycentricWeights()
    :param srcWeightDict: <dict> clusterWeight dict structured by eg. the Skinning module's .getSkinWeights()
    :return: <dict> of final cluster weights for the transfer
    '''

    zippedCluster = {}

    for tgtVtxId in baryWeightDict.keys():
        zippedCluster[tgtVtxId] = {}

        for srcVtxId, baryWeight in baryWeightDict[tgtVtxId].items():

            for infVtxId, infWeight in clusterWeightDict[srcVtxId].items():
                zippedWeight = infWeight * baryWeight

                # if influence already exists, add them together
                if infVtxId in zippedCluster[tgtVtxId].keys():
                    zippedWeight += zippedCluster[tgtVtxId][infVtxId]

                zippedCluster[tgtVtxId][infVtxId] = zippedWeight

    return zippedCluster


def getVtxCount(shapeNode):
    '''
    Gets the vertex count for a given shape node
    :param shapeNode:
    :return: <int> vertex count
    '''
    # lazy buffer if object is passed
    shapeNode = getShape(shapeNode)
    return len(shapeNode.vtx)


def filePathPrompt(dialogMode, caption="SkinWeights", dirPath="D:/", filter="API Skin weights file (*.skinWeights)"):
    '''
    Generates a maya file dialogue window and returns a filepath
    :param fileMode: <int> 0 is export path, 1 is save path
    :param caption: <string> title of dialog
    :param dirPath: <string> starting directory
    :param filter: <string> file type filter in format "(*.fileType)"
    :return: <string> file path
    '''

    filePath = pm.system.fileDialog2(fileMode=dialogMode,
                                     caption=caption,
                                     dir=dirPath,
                                     fileFilter=filter)
    if filePath:
        # fileDialog2 returns a list if successful so we fetch the first element
        filePath = filePath[0]

    return filePath


def matchDictionaryToSceneMeshes(weightDictionary, selected=False):
    '''
    Matches names of the dictionary.keys() with meshes in the scene.
    Matches ALL objects, then filters based on user selection
    :param weightDictionary: <dict> weightDictionary generated by the apiVtxAttribs module
    :param selected: <bool> operate on selection
    :return: <dict> <list>, reformatted weight dictionary, found matching objects in scene
    '''
    validNodeList = []
    sceneWeightDict = {}
    # find all matching or valid objects in the scene
    for dictNodeName in weightDictionary.keys():
        sceneNodeList = pm.ls(dictNodeName, recursive=True, type="transform")

        if not len(sceneNodeList):
            # we have not found matches by the stored name, try short names
            shortDictNodeName = dictNodeName.split("|")[-1]
            sceneNodeList = pm.ls(shortDictNodeName, recursive=True, type="transform")
        print(sceneNodeList)
        for sceneNode in sceneNodeList:
            if vtxCountMatch(sceneNode, weightDictionary[dictNodeName]["vtxCount"]):
                # found match on both name and vtxcount, copy info to the local sceneWeightDict
                sceneWeightDict[sceneNode.name()] = weightDictionary[dictNodeName]
                validNodeList.append(sceneNode.name())

    # filter on selection
    if selected:
        selectionMatchedList = []
        selectedNodes = getTransforms(pm.ls(sl=True))
        for selectedNode in selectedNodes:
            if selectedNode in validNodeList:
                selectionMatchedList.append(str(selectedNode))

        # replace validNodeList with selection match
        validNodeList = selectionMatchedList

    if sceneWeightDict and validNodeList:
        return sceneWeightDict, validNodeList
    else:
        return {}, []


def vtxCountMatch(node, vtxCount):
    '''
    Queries if the vtx count of the given node and the given vtxCount matches
    :param node: <shape> or <transform>
    :param vtxCount: <int>
    :return: <bool>
    '''
    node = getShape(node)
    try:
        nodeVtxCount = len(node.vtx)
        if nodeVtxCount == vtxCount:
            return True
        return False
    except:
        print("Unable to retrieve .vtx list from %s" % node)


def getPickleObject(filePath):
    '''
    Unpacks a the pickled weights files into a human readable weight dictionary
    :param filePath: <string>
    :return: <dict>
    '''
    with open(filePath, 'rb') as fp:
        try:
            pickleObject = pickle.Unpickler(fp)
            return pickleObject.load()
        except:
            raise Exception("Unable to unpack %s" % filePath)


def pickleDumpWeightsToFile(weightDict, filePath):
    '''
    Export a weights file and handle versioncontrol
    :param weightDict: <dict> Dictionary of clusterweights
    :param filePath: <string>
    :return:
    '''

    # if vc.checkout([filePath]):
    #     with open(filePath, 'wb') as fp:
    #         pickle.dump(weightDict, fp, pickle.HIGHEST_PROTOCOL)
    #     vc.add([filePath])
    # else:
    #     with open(filePath, 'wb') as fp:
    #         pickle.dump(weightDict, fp, pickle.HIGHEST_PROTOCOL)
    #     print "WARNING: Successfully exported weights but was unable to connect to perforce"

    with open(filePath, 'wb') as fp:
        pickle.dump(weightDict, fp, pickle.HIGHEST_PROTOCOL)


def handleTransferNodesList(transferNodes=None):
    '''
    Handle empty transferNodeslist in transfer functions
    Basically fetches the maya scene selection to the transfernodeslist if nothing is provided
    Checks if all transfer objects are valid mesh objects
    :param transferNodes: <list>
    :return: <list>
    '''

    # operate on selection
    if transferNodes is None:
        transferNodes = getTransforms(pm.ls(os=True))

    if len(transferNodes) < 2:
        print("# Error: Select atleast 2 objects, a source and a target")
        return []

    if None in getShapes(transferNodes):
        print("# Error: Non-mesh objects found in selection")
        return []

    return transferNodes


class LoadingBar():

    def __init__(self):
        self.lprogressAmount = 0

    def loadingBar(self, ltitle, lmaxValue, lstatus="Working..."):
        #if loadingBar is accessed the first time, create progressWindow
        if self.lprogressAmount == 0:
            pm.progressWindow(title=ltitle, progress=self.lprogressAmount, status=lstatus, isInterruptable=False, max=lmaxValue)
        self.lprogressAmount += 1
        #increases progress in progressWindow
        pm.progressWindow(edit=True, progress=self.lprogressAmount)
        #if progressAmount is lmaxValue, finish progressWindow
        if self.lprogressAmount == lmaxValue:
            pm.progressWindow(endProgress=True)
            self.progressAmount = 0