AlignVerts_UI.py

import copy
import collections
import traceback

from maya import cmds
import pymel.core as pm



class AlignVerts():
    
    # for fancy coloring of buttons
    COLOR_RED = [0.8, 0.5, 0.5]
    COLOR_GREEN = [0.5, 0.8, 0.5]
    
    
    def __init__(self):
        
        self.vert_collection_1 = collections.OrderedDict()
        self.vert_collection_2 = collections.OrderedDict()
        self.big_button_was_pressed = False
        self.undo_has_started = False
        
        
        window_title = "Align Verts"
        window_internal_name = window_title.replace(" ", "_")
        
        if pm.window(window_internal_name, exists=True):
            pm.deleteUI(window_internal_name)
        
        with pm.window(window_internal_name, title=window_title) as win:
            with pm.verticalLayout() as main_layout:
                    
                
                with pm.horizontalLayout() as vert_col_1_layout:
                    self.vert_col_1_btn = pm.button(label="Set Vert Collection 1", command=self.cache_verts_1)
                    self.vert_col_1_btn.setBackgroundColor(self.COLOR_RED)
                    self.move_verts_1 = pm.checkBox("Move", value=True, changeCommand=self.launch_command)
                vert_col_1_layout.redistribute(2,1)
                
                with pm.horizontalLayout():
                    pm.text("Move Axis: ")
                    self.col_1 = pm.radioCollection()
                    pm.radioButton(label="X", changeCommand=self.launch_command, select=True)
                    pm.radioButton(label="Y", changeCommand=self.launch_command)
                    pm.radioButton(label="Z", changeCommand=self.launch_command)
                
                
                pm.separator()
                
                with pm.horizontalLayout() as vert_col_2_layout:
                    self.vert_col_2_btn = pm.button(label="Set Vert Collection 2", command=self.cache_verts_2)
                    self.vert_col_2_btn.setBackgroundColor(self.COLOR_RED)
                    self.move_verts_2 = pm.checkBox("Move", value=True, changeCommand=self.launch_command)
                vert_col_2_layout.redistribute(2,1)
                
                
                with pm.horizontalLayout():
                    pm.text("Move Axis: ")
                    
                    self.col_2 = pm.radioCollection()
                    
                    pm.radioButton(label="X", changeCommand=self.launch_command)
                    pm.radioButton(label="Y", changeCommand=self.launch_command)
                    pm.radioButton(label="Z", changeCommand=self.launch_command, select=True)
                    
                
                pm.separator()
                
                with pm.horizontalLayout():
                    pm.text("Shift Vertex List")
                    self.offset_int = pm.intField("Offset", step=1, value=0, maxValue=10000, minValue=-10000, 
                                                  changeCommand=self.launch_command,
                                                  dragCommand=lambda x: self.launch_command(dragging=True))# don't end undo while the value is changing
                
                with pm.verticalLayout() as btn_layout:
                    pm.button("Align", command=lambda x: self.launch_command(big_button=True))
                    pm.button("Reset", command=self.reset_tool)
                btn_layout.redistribute(2,1)
                
            main_layout.redistribute(4,4,1,4,4,1,4,10)
                
    
    def cache_verts_1(self, *args):
        sel_verts = pm.selected(flatten=True)
        if not sel_verts:
            pm.warning("No vertices selected")
            return
            
        self.vert_collection_1.clear()
        for vert in sel_verts:
            self.vert_collection_1[vert] = list(vert.getPosition(space="world"))
            
        self.vert_col_1_btn.setBackgroundColor(self.COLOR_GREEN)
    
    def cache_verts_2(self, *args):
        sel_verts = pm.selected(flatten=True)
        if not sel_verts:
            pm.warning("No vertices selected")
            return
        
        self.vert_collection_2.clear()
        for vert in sel_verts:
            self.vert_collection_2[vert] = list(vert.getPosition(space="world"))
            
        self.vert_col_2_btn.setBackgroundColor(self.COLOR_GREEN)
        
    def reset_tool(self, *args):
        self.reset_mesh()
        self.big_button_was_pressed = False
    
    def reset_mesh(self, *args):
        """
        Brings back the old vert positions
        """
        [vert.setPosition(vert_pos, space="world") for vert, vert_pos in self.vert_collection_1.items()]
        [vert.setPosition(vert_pos, space="world") for vert, vert_pos in self.vert_collection_2.items()]
        
    
    def launch_command(self, arg_catcher=(), big_button=False, dragging=False):
        if big_button:
            self.big_button_was_pressed = True
        
        if not self.big_button_was_pressed:
            return
            
        
        if not self.vert_collection_1.keys() or not self.vert_collection_2.keys():
            pm.warning("Vert Collections need to be defined")
            return
        
        
        if not dragging: # skip on dragging for performance
            self.reset_mesh() # reset mesh so when we toggle different modes that we always start with the same thing
        
        
        # start undo chunk for the offset int field
        if dragging and not self.undo_has_started:
            cmds.undoInfo(openChunk=True)
            self.undo_has_started = True
        
            
        try:
            # get values from UI
            axis_dict = {"X":0, "Y": 1, "Z":2}
            col_1_axis_index = axis_dict[pm.radioButton(self.col_1.getSelect(), q=True, label=True)]
            col_2_axis_index = axis_dict[pm.radioButton(self.col_2.getSelect(), q=True, label=True)]
            
            offset = self.offset_int.getValue()
            move_col_1 = self.move_verts_1.getValue()
            move_col_2 = self.move_verts_2.getValue()
            
            
            
            # Actual Logic
            rotated_verts_2 = collections.deque(self.vert_collection_2.keys())
            rotated_verts_2.rotate(offset) # shift the vertex list by an amount
            
            # The vert position lists are mutable, so we want to deepcopy these to make sure we don't modify the reset positions
            copied_collection_1 = copy.deepcopy(self.vert_collection_1)
            copied_collection_2 = copy.deepcopy(self.vert_collection_2)
            
            
            for vert_1, vert_2 in zip(self.vert_collection_1.keys(), rotated_verts_2): # will end on whichever list is shorter
                vert_1_pos = copied_collection_1[vert_1]
                vert_2_pos = copied_collection_2[vert_2]
                
                if move_col_1:
                    vert_1_pos[col_1_axis_index] = vert_2_pos[col_1_axis_index]
                    vert_1.setPosition(vert_1_pos, space="world")
                    
                if move_col_2:
                    vert_2_pos[col_2_axis_index] = vert_1_pos[col_2_axis_index]    
                    vert_2.setPosition(vert_2_pos, space="world")
                
        except StandardError as e: # try excepting this to try and prevent breaking the undo queue
            traceback.print_exc()
        
        # close undo chunk
        if not dragging and self.undo_has_started:
            cmds.undoInfo(closeChunk=True)
            self.undo_has_started = False



if __name__ == "__main__":
    win = AlignVerts()