Untitled

from maya.api import OpenMaya as om
from maya.api import OpenMayaAnim as oma
import sys


def maya_useNewAPI():
    """
    The presence of this function tells Maya that the plugin produces, and
    expects to be passed, objects created using the Maya Python API 2.0.
    """
    pass


k_help_flag = '-h'
k_help_long_flag = '-help'

help_message = 'This command smooths skinCluster weights while respecting the lockInfluenceWeights attribute value '\
               'of the influences'


class SmoothSkinClusterWeights(om.MPxCommand):

    K_PLUGIN_CMD_NAME = 'smoothSkinClusterWeights'

    sel = None

    # used for undo
    skin_objs = []
    skin_weights_original = []
    mesh_dags = []
    inf_ias = []

    def __init__(self):

        super(SmoothSkinClusterWeights, self).__init__()

    def doIt(self, args):

        syntax = self.syntax()
        try:
            parsed_arguments = om.MArgDatabase(syntax, args)
            self.sel = parsed_arguments.getObjectList()
        except:
            om.MGlobal.displayError('Error getting objects')
            return

        if parsed_arguments.isFlagSet(k_help_flag):
            self.setResult(help_message)
        if parsed_arguments.isFlagSet(k_help_long_flag):
            self.setResult(help_message)

        if self.sel.length() > 0:
            self.redoIt()

    def redoIt(self):

        self.skin_objs = []
        self.skin_weights_original = []
        self.mesh_dags = []
        self.inf_ias = []

        unsuccessful = []

        for i in range(self.sel.length()):

            dep = self.sel.getDependNode(i)
            dag = None
            comp = None

            if dep.hasFn(om.MFn.kTransform):

                dag, comp = self.sel.getComponent(i)

                dag_fn = om.MFnDagNode(dag)
                num_children = dag_fn.childCount()
                for j in range(num_children):
                    child_obj = dag_fn.child(j)
                    if child_obj.hasFn(om.MFn.kMesh):
                        dep = child_obj

                        dag_fn.setObject(child_obj)
                        temp_sel = om.MSelectionList()
                        temp_sel.add(dag_fn.fullPathName())
                        dag, comp = temp_sel.getComponent(0)
                        break

            skin_obj = None

            if dep.hasFn(om.MFn.kMesh):

                if dag is None and comp is None:
                    dag, comp = self.sel.getComponent(i)

                dep_it = om.MItDependencyGraph(dep, filter=om.MFn.kSkinClusterFilter, direction=om.MItDependencyGraph.kUpstream)
                while not dep_it.isDone():
                    skin_obj = dep_it.currentNode()
                    break

            if dep.hasFn(om.MFn.kSkinClusterFilter):

                skin_obj = dep

                dep_it = om.MItDependencyGraph(dep, filter=om.MFn.kMesh, direction=om.MItDependencyGraph.kUpstream)
                while not dep_it.isDone():

                    current_node = dep_it.currentNode()

                    temp_dep_fn = om.MFnDependencyNode(current_node)
                    temp_sel = om.MSelectionList()
                    temp_sel.add(temp_dep_fn.uniqueName())
                    dag, comp = temp_sel.getComponent(0)
                    break

            elif not dag.hasFn(om.MFn.kMesh):
                unsuccessful.append((dag, 'mesh shape cannot be found'))
                continue

            if not comp.isNull() and not comp.hasFn(om.MFn.kMeshVertComponent):
                unsuccessful.append((dag, 'component selection is not vertex type'))
                continue

            if skin_obj is None:
                unsuccessful.append((dag, 'has no skinCluster deformer'))
                continue

            skin_fn = oma.MFnSkinCluster(skin_obj)
            inf_dpa = skin_fn.influenceObjects()
            num_infs = len(inf_dpa)

            if num_infs < 2:
                unsuccessful.append((dag, 'needs 2 or more skinCluster influences for smoothing'))
                continue

            inf_ia = om.MIntArray(range(num_infs))

            # get the skin weights
            skin_wts = skin_fn.getWeights(dag, om.MObject(), inf_ia)
            skin_wts_original = skin_fn.getWeights(dag, om.MObject(), inf_ia)

            self.mesh_dags.append(dag)
            self.skin_objs.append(skin_obj)
            self.inf_ias.append(inf_ia)
            self.skin_weights_original.append(skin_wts_original)

            # get lockInfluenceWeights values of influences
            inf_liws = []
            for inf_dp in inf_dpa:
                inf_fn_dag = om.MFnDagNode(inf_dp)
                inf_liw_plug = inf_fn_dag.findPlug('lockInfluenceWeights', False)
                inf_liws.append(inf_liw_plug.asBool())

            mesh_it = om.MItMeshVertex(dag)
            mesh_comp_it = om.MItMeshVertex(dag, comp)

            # get weights for each vertex
            vertices_wts = []
            while not mesh_it.isDone():
                index = mesh_it.index()
                vertices_wts.append(skin_wts[index * num_infs: (index + 1) * num_infs])

                mesh_it.next()

            mesh_it.reset()

            while not mesh_comp_it.isDone():

                index = mesh_comp_it.index()
                mesh_it.setIndex(index)

                vert_wts = vertices_wts[index]

                connected_indices = mesh_it.getConnectedVertices()
                connected_wts = [vertices_wts[j] for j in connected_indices]

                blend_wts = []

                for j in range(num_infs):
                    vert_inf_wt = vert_wts[j]
                    connected_inf_wts = sum([connected_inf_wt[j] for connected_inf_wt in connected_wts])
                    blend_inf_wt = (vert_inf_wt + connected_inf_wts / len(connected_wts)) / 2

                    blend_wts.append(blend_inf_wt)

                if any(inf_liws):

                    locked_wt = 0
                    unlocked_wt = 0

                    for j, inf_liw in enumerate(inf_liws):

                        if inf_liw:
                            locked_wt += vert_wts[j]
                            blend_wts[j] = vert_wts[j]
                        else:
                            unlocked_wt += blend_wts[j]

                    if unlocked_wt != 0:
                        for j, inf_liw in enumerate(inf_liws):
                            if not inf_liw:
                                blend_wts[j] *= (1 - locked_wt) / unlocked_wt

                for j in range(num_infs):
                    skin_wts[num_infs * index + j] = blend_wts[j]

                mesh_comp_it.next()

            skin_fn.setWeights(dag, om.MObject(), inf_ia, skin_wts)

        if unsuccessful:
            for dag, message in unsuccessful:
                om.MGlobal.displayWarning(f'{dag.partialPathName()} {message}')

    def undoIt(self):

        for skin_obj, _skin_weights, dag, inf_ia in zip(
                self.skin_objs, self.skin_weights_original, self.mesh_dags, self.inf_ias):
            skin_fn = oma.MFnSkinCluster(skin_obj)
            skin_fn.setWeights(dag, om.MObject(), inf_ia, _skin_weights)

    def isUndoable(self):

        return True

    @staticmethod
    def cmdCreator():

        return SmoothSkinClusterWeights()

    @staticmethod
    def syntaxCreator():

        syntax = om.MSyntax()
        syntax.useSelectionAsDefault(True)
        syntax.setObjectType(syntax.kSelectionList)
        syntax.setMinObjects(0)

        syntax.addFlag(k_help_flag, k_help_long_flag)

        return syntax


def initializePlugin(plugin):
    plugin_fn = om.MFnPlugin(plugin)

    try:
        plugin_fn.registerCommand(
            SmoothSkinClusterWeights.K_PLUGIN_CMD_NAME, SmoothSkinClusterWeights.cmdCreator, SmoothSkinClusterWeights.syntaxCreator)
    except:
        sys.stderr.write(f'Failed to register command: {SmoothSkinClusterWeights.K_PLUGIN_CMD_NAME}')


def uninitializePlugin(plugin):
    plugin_fn = om.MFnPlugin(plugin)

    try:
        plugin_fn.deregisterCommand(SmoothSkinClusterWeights.K_PLUGIN_CMD_NAME)
    except:
        sys.stderr.write(f'Failed to unregister command: {SmoothSkinClusterWeights.K_PLUGIN_CMD_NAME}')