# ----------------------------------------------------------------
# Creation Info:
# Author: Sunny Chopra http://underearth.wordpress.com
# Thanks to Simon for Prototype of this code and guiding and inspiring me to start diving to maya API
# Date: 19.04.2009
# Version: 0.3
# 
# Release Notes:
# 0.1- 19.04.2009 - Initial release
# 0.2- 19.04.2009 - Undo state implemented 
# 0.3- 20.01.2012 - some optimization added vertex ID
# 0.4- 07.02.2012 - fixed bugs in importing weights
# 0.45- 19.02.2012 - fixed more bugs in importing weights, now works on component mode

# This program is free software; you can redistribute it
# and/or modify it under the terms of the GNU General
# Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.



""" Usage/Options:
Select skin mesh and run in script editor

scSaveWeights -roundOff 5 -action "export" -file "d:/weights.txt"
scSaveWeights -action "import" -file "d:/weights.txt"

Flags:
roundOff / r = (int)Decimal Position to roundOff to.
action / a = (string) Either Export of Import to save or import skinWeights
file/ f = (string) File path to save file into, eg "c:/weights.w"
quick / qi = (string) For importing/exporting out skinWeights in as it is mode, which is faster, doesn't work on components

 To dos:
 Save positional data with export and options for importing with world Space pos.
 Some error check if number of vertices mismatch.
 Mirroring skinWeights.
 Search and replace names for joint.
"""

import maya.OpenMaya as om
import maya.OpenMayaAnim as omAnim
import maya.OpenMayaMPx as OpenMayaMPx
import maya.cmds as mc
import sys

class statusError(Exception):
    pass

kPluginCmdName = "scSaveWeights" # name of the command

kQuickSaveFlag = "-qi"
kQuickSaveLongFlag = "-quick"
kRoundOffFlag = "-r"
kRoundOffLongFlag = "-roundOff"
kActionFlag = "-a"
kActionLongFlag = "-action"
kFileNameFlag = "-f"
kFileNameLongFlag = "-file"

# Define the command

########################################################################
class scSaveWeights(OpenMayaMPx.MPxCommand):
    """class for exporting out maya skincluster into Ascii file and reading them into geometry
    currently supporting Polygon geo's only"""

    #----------------------------------------------------------------------
    def __init__(self):
        """Constructor"""

        OpenMayaMPx.MPxCommand.__init__(self)

        #default value for rounding off save weights,setup private data members
        self.__round = 4
        self.__defaultFileName = "c:/weightFile.txt"

    #----------------------------------------------------------------------
    def doIt(self,args):
        """collecting all information to do real stuff"""
        argData = om.MArgDatabase(self.syntax(), args)
        if argData.isFlagSet(kQuickSaveFlag):
            self.__QuickSave = True
        else:
            self.__QuickSave = False
        if argData.isFlagSet(kRoundOffFlag):
            self.__round = argData.flagArgumentInt(kRoundOffFlag, 0)
        else:
            self.__round = self.__round
        if not argData.isFlagSet(kActionFlag):
            raise statusError("action flag must be set with options of either export/import")
        else:
            self.__action = argData.flagArgumentString(kActionFlag, 0)
        if argData.isFlagSet(kFileNameFlag):
            self.__fileName = argData.flagArgumentString(kFileNameFlag, 0)
        else:
            self.__fileName = self.__defaultFileName

        self.redoIt()

    #----------------------------------------------------------------------
    def redoIt(self):
        """Now for doing real stuff"""

        selList  = om.MSelectionList()
        self.__dagPathGeo = om.MDagPath()
        self.__indArrayImport= om.MIntArray()
        self.__allVert = om.MObject() # shall put selected vertices into it , create selection list for components
        
        om.MGlobal.getActiveSelectionList(selList)
        if selList.length() == 0:
            raise statusError("Nothing selected")
        elif selList.length() > 1:
            raise statusError("Too many objects selected")

        selOrig = mc.ls(sl=1)
        selStrComp = mc.polyListComponentConversion (ff=1,fe=1,fuv=1,fvf=1,fv=1,tv=1) #ADD vertex to sel list
        mc.select(selStrComp,r=1)
        om.MGlobal.getActiveSelectionList(selList)
        #om.MGlobal.getSelectionListByName( selStrComp,selList)
        mc.select(selOrig,r=1)

        selList.getDagPath(0,self.__dagPathGeo,self.__allVert)

        if self.__dagPathGeo.apiType() != 295 and self.__dagPathGeo.apiType() != 294 and self.__dagPathGeo.apiType() != 110:
            raise statusError("Please select Polygon object, or vertices to save weights")

        # approach to skinclusterNode
        for eachNode in mc.listHistory(self.__dagPathGeo.partialPathName(),f=0,bf=1):
            if mc.nodeType(eachNode) == 'skinCluster':
                break
        om.MGlobal.getSelectionListByName( eachNode,selList)
        skinObj = om.MObject()
        selList.getDependNode(1, skinObj )
        self.__skinClusterNode = omAnim.MFnSkinCluster(skinObj)
        skinObj.apiTypeStr()
        if self.__action == "export":
            self.__scSaveWeightsExport()
        elif self.__action == "import":
            self.__scSaveWeightsImport()

    #----------------------------------------------------------------------
    def undoIt(self):
        if self.__action == "import":
            self.__skinClusterNode.setWeights(self.__dagPathGeo,self.__allVert,self.__indArrayUndo,self.__unDoWeights,False)
            # i used UNDO array(of JOINT INDEX) which i created while zeroing out all weights in import function

    def isUndoable(self):
        return True

    #----------------------------------------------------------------------
    def __scSaveWeightsExport(self):
        """will write out file to specified location"""

        # skincluster export
        weights = om.MDoubleArray()
        #componentIter = om.MItGeometry(self.__dagPathGeo) # does iterate on whole mesh
        componentIter = om.MItGeometry(self.__dagPathGeo,self.__allVert) # iterate only on selected components


        # creating pointer objetc for to be used in MFnSkinCluster's get weight function
        infCount = om.MScriptUtil()
        infCountPtr = infCount.asUintPtr()
        om.MScriptUtil.setUint(infCountPtr, 0)

        # writes out skin weights to ascii format.
        f = open(self.__fileName,'wb')
        f.write("#starting of skinWeight data\n")

        infNameList =[]
        infNameList = self.__scCreateInfNameList(self.__skinClusterNode)
        toWriteInfs =  "# Joint list\n"
        toWriteInfs += " ".join(infNameList)
        toWriteInfs += "\n# Joint list ends\n"
        f.write(toWriteInfs)
        f.write("# skin weight begins\n")
        if self.__QuickSave == 1:
            toWriteLineS =""
            self.__skinClusterNode.getWeights(self.__dagPathGeo,self.__allVert,weights,infCountPtr)

            for i in range(0,len(weights),1):
                toWriteLineS += str(weights[i])+" "
            toWriteLineS += "\n"
            f.write(toWriteLineS)
            toWriteLineS=""

        else:
            while componentIter.isDone() == 0:

                self.__skinClusterNode.getWeights(self.__dagPathGeo, componentIter.currentItem (), weights, infCountPtr)
                weightDict = self.__scNormalizeWeight(weights,self.__round,infNameList,0)
                toWriteLineS = "".join('['+`componentIter.index()`+']')
                toWriteLineS += "\n"
                toWriteLineS += " ".join(weightDict.keys())
                toWriteLineS += "\n"
                for eachVal in weightDict.values():
                    toWriteLineS += str(eachVal) + " "
                toWriteLineS += "\n"
                f.write(toWriteLineS)
                componentIter.next()

        f.write("//end of skinWeight data")
        f.close()

    #----------------------------------------------------------------------
    def __scSaveWeightsImport(self):
        """will read in weights to geometry from given file"""
        eachVertImport = om.MObject()
        weights = om.MDoubleArray()

        self.__indArrayUndo= om.MIntArray()
        self.__unDoWeights = om.MDoubleArray()

        infNameList=[]
        infNameList = self.__scCreateInfNameList(self.__skinClusterNode)

        for i in range(0,len(infNameList),1):
            self.__indArrayUndo.append(i)
            weights.append(0)
        #zero out all the weights to supress error message
        self.__skinClusterNode.setWeights(self.__dagPathGeo,self.__allVert,self.__indArrayUndo,weights,False,self.__unDoWeights)
        weights.clear()
        self.__indArrayImport.clear()
        # start reading the file
        fRead = open(self.__fileName,'rb')
        for line in fRead:
            if line.find('# skin weight begins') != -1:
                break
        
        if self.__QuickSave == 1:
            #weightsString =""
            weightsString = fRead.next().split()
            for i in range(len(weightsString)):
                weights.append(float(weightsString[i]))
            for i in range(0,len(infNameList),1):
                self.__indArrayImport.append(i)

            self.__skinClusterNode.setWeights(self.__dagPathGeo,self.__allVert,self.__indArrayImport,weights,False)
            weights.clear()
            self.__indArrayImport.clear()
            fRead.close()
            weightsString =""
        elif self.__QuickSave == 0:
        #read out weights file
            allVtxDict = {}
            for line in fRead:
                if line.find("//end of skinWeight data") != -1: break
                dictVtxKey = line[1:-2] # removes [] #make sure data [] matches correctly at this point
                ValJntList = fRead.next().split()
                ValWtList = fRead.next().split()
                allVtxDict[dictVtxKey] = [ValJntList,ValWtList]
            fRead.close()
            componentIter = om.MItGeometry(self.__dagPathGeo,self.__allVert) # iterate only on selected components
            while componentIter.isDone() == 0:
                vNumb = componentIter.index()
                infName = allVtxDict[`vNumb`][0]
                weightsString = allVtxDict[`vNumb`][1]
                for i in range(len(weightsString)):
                    weights.append(float(weightsString[i]))
                    self.__indArrayImport.append(infNameList.index(infName[i]))

                self.__skinClusterNode.setWeights(self.__dagPathGeo,componentIter.currentItem(),self.__indArrayImport,weights,False)
                componentIter.next()
                weights.clear()
                self.__indArrayImport.clear()

    #----------------------------------------------------------------------
    def __scNormalizeWeight(self,weights,place,infNameList,addZero):
        """For the given weight list this func will normalize and round off the weights and
        provisionally could remove nonZero members.This func returns Dictionary with key as
        influenceObj and value as weights."""

        allWeight = 0
        largestWeight = 0
        heavyJnt ="" 
        nonZeroDict={}
        for i in range(0,len(weights),1):
            if (weights[i] != 0):
                weights[i]= round(weights[i],4)
            nonZeroDict[infNameList[i]] = weights[i]
            if (largestWeight < nonZeroDict[infNameList[i]]):
                largestWeight = nonZeroDict[infNameList[i]]
                heavyJnt = infNameList[i]
            allWeight += nonZeroDict[infNameList[i]]
        remWeight = 1 - allWeight
        nonZeroDict[heavyJnt] = nonZeroDict[heavyJnt] + remWeight
        if addZero == 0:
            for each in nonZeroDict.keys():
                if nonZeroDict[each] == 0:
                    nonZeroDict.pop(each)
        return nonZeroDict

    #----------------------------------------------------------------------
    def __scCreateInfNameList(self,skinClusterNode):
        """this function return out list containing joint name in ordered list as they are attached internally to skinCluster
        which are used when seting or getting weights from skincluster."""

        infNameList =[]
        infs = om.MDagPathArray()
        numInfs = skinClusterNode.influenceObjects(infs)
        for counter in range(0,numInfs,1):
            infName = infs[counter].partialPathName()
            infNameList.append(infName)
        return infNameList

# code for maya to understand this as scripted plugin
def cmdCreator():
    return OpenMayaMPx.asMPxPtr(scSaveWeights())

def syntaxCreator():
    syntax = om.MSyntax()
    syntax.addFlag(kQuickSaveFlag,kQuickSaveLongFlag,om.MSyntax.kNoArg)
    syntax.addFlag(kRoundOffFlag, kRoundOffLongFlag,om.MSyntax.kLong)
    syntax.addFlag(kActionFlag, kActionLongFlag, om.MSyntax.kString)
    syntax.addFlag(kFileNameFlag, kFileNameLongFlag, om.MSyntax.kString)
    return syntax

def initializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject,"sunnyC",".2","Any")
    try:
        mplugin.registerCommand( kPluginCmdName, cmdCreator, syntaxCreator )
    except:
        raise statusError( "Failed to register command: %s\n" % kPluginCmdName )


def uninitializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject)
    try:
        mplugin.deregisterCommand( kPluginCmdName )
    except:
        raise statusError( "Failed to unregister command: %s\n" % kPluginCmdName )
        


# code optimaztion
# while loop is eveil also iterate using this, please check if dictionariy is also creating memory problem, than also use simple iterate kill dict