#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
版权所有: 清泉时代科技有限公司
联系方式: q.100@qq.com
创建日期: 2023/08/08
"""

import maya.cmds as cmds
import maya.mel as mel

def sg_get_orig_name(geo):
    """
    获取模型的Orig节点名称列表
    
    参数:
        geo (str): 模型名称
    返回:
        list: Orig节点名称列表
    """
    orig = []
    if not geo or not cmds.objExists(geo):
        return orig
    
    shape_mesh = cmds.listRelatives(geo, shapes=True) or []
    for shape in shape_mesh:
        if "Orig" in shape:
            orig.append(shape)
    return orig

def sg_get_orig_shape_name(geo):
    """
    获取模型的带有groupParts连接的Orig节点名称列表
    
    参数:
        geo (str): 模型名称
    返回:
        list: 符合条件的Orig节点名称列表
    """
    orig = []
    if not geo or not cmds.objExists(geo):
        return orig
    
    shape_mesh = cmds.listRelatives(geo, shapes=True) or []
    for shape in shape_mesh:
        if "Orig" in shape:
            connections = cmds.listConnections(shape) or []
            for conn in connections:
                if "groupParts" in conn:
                    orig.append(shape)
                    break
    return orig

def sg_bind_pose_reset(objs):
    """
    重置模型的绑定姿势
    
    参数:
        objs (list): 要处理的模型列表
    """
    # 开始进度条
    cmds.SGProgressBar(sp=True)
    cmds.SGProgressBar(max=len(objs))
    
    for obj in objs:
        # 更新进度条
        cmds.SGProgressBar(t=f"[{obj}] Bind Pose Reset...")
        cmds.SGProgressBar(apr=1)
        
        # 获取蒙皮变形节点
        skin_cluster = mel.eval(f'findRelatedSkinCluster("{obj}")')
        
        if cmds.objExists(skin_cluster):
            # 获取混合变形节点
            blend_shape = cmds.SGGetBlendShape(obj)
            
            # 创建临时复制模型
            copy_name = cmds.duplicate(obj, name=f"{obj}_CopyTarget", returnRootsOnly=True)
            cmds.blendShape(copy_name[0])
            cmds.delete(copy_name[0], constructionHistory=True)
            
            # 获取骨骼并应用蒙皮
            bones = cmds.skinCluster(skin_cluster, query=True, influence=True)
            cmds.skinCluster(bones, copy_name[0], toSelectedBones=True)
            
            # 复制蒙皮权重
            cmds.copySkinWeights(
                sourceShape=obj,
                destinationShape=copy_name[0],
                noMirror=True,
                surfaceAssociation="closestPoint",
                influenceAssociation=["closestJoint", "oneToOne"]
            )
            
            # 解绑原模型
            cmds.skinCluster(obj, edit=True, unbind=True)
            
            # 获取复制模型的形状节点
            target_shape = cmds.pickWalk(copy_name[0], direction="down")
            
            if cmds.objExists(blend_shape):
                # 处理带有混合变形的情况
                orig = sg_get_orig_name(obj)
                try:
                    cmds.connectAttr(f"{orig[0]}.outMesh", f"{blend_shape}.originalGeometry[0]", force=True)
                    cmds.connectAttr(f"{orig[0]}.worldMesh[0]", f"{blend_shape}.input[0].inputGeometry", force=True)
                except: pass
                
                cmds.currentTime(0)
                cmds.connectAttr(f"{target_shape[0]}.outMesh", f"{orig[0]}.inMesh", force=True)
                cmds.currentTime(1)
                cmds.disconnectAttr(f"{target_shape[0]}.outMesh", f"{orig[0]}.inMesh")
                cmds.currentTime(0)
            else:
                # 处理不带混合变形的情况
                orig = sg_get_orig_shape_name(obj)
                if orig:
                    cmds.currentTime(0)
                    cmds.connectAttr(f"{target_shape[0]}.outMesh", f"{orig[0]}.inMesh", force=True)
                    cmds.currentTime(1)
                    cmds.disconnectAttr(f"{target_shape[0]}.outMesh", f"{orig[0]}.inMesh")
                    cmds.currentTime(0)
                    
                    # 重置控制点位置
                    vertex_count = cmds.polyEvaluate(orig[0], vertex=True)
                    for i in range(vertex_count):
                        cmds.setAttr(f"{orig[0]}.controlPoints[{i}]", 0, 0, 0, type="float3")
                else:
                    # 处理没有Orig节点的情况
                    still_shape = cmds.pickWalk(obj, direction="down")
                    cmds.currentTime(0)
                    cmds.connectAttr(f"{target_shape[0]}.outMesh", f"{still_shape[0]}.inMesh", force=True)
                    cmds.currentTime(1)
                    cmds.disconnectAttr(f"{target_shape[0]}.outMesh", f"{still_shape[0]}.inMesh")
                    cmds.currentTime(0)
            
            # 重新绑定蒙皮并复制权重
            cmds.skinCluster(bones, obj, name=skin_cluster, toSelectedBones=True)
            cmds.copySkinWeights(
                sourceShape=copy_name[0],
                destinationShape=obj,
                noMirror=True,
                surfaceAssociation="closestPoint",
                influenceAssociation=["closestJoint", "oneToOne"]
            )
            
            # 删除临时模型
            cmds.delete(copy_name[0])
    
    # 结束进度条
    cmds.SGProgressBar(ep=True)