import maya.cmds as cmds
import maya.OpenMayaUI as omui
import maya.OpenMaya as om
from PySide2 import QtWidgets, QtCore
from shiboken2 import wrapInstance

class VertexDistributor(QtWidgets.QWidget):
    def __init__(self):
        super(VertexDistributor, self).__init__()
        self.setWindowTitle("Vertex Distributor")
        self.setWindowFlags(self.windowFlags() | QtCore.Qt.WindowStaysOnTopHint)
        self.setGeometry(100, 100, 500, 200)
        self.create_ui()

    def create_ui(self):
        layout = QtWidgets.QVBoxLayout()

        # Create a form layout for aligned fields
        form_layout = QtWidgets.QFormLayout()

        # Start Point
        start_layout = QtWidgets.QHBoxLayout()
        self.start_field = QtWidgets.QLineEdit()
        start_layout.addWidget(self.start_field)
        start_layout.addWidget(self.create_button("Assign", self.assign_start))
        start_layout.addWidget(self.create_button("Select", self.select_start))
        start_layout.addWidget(self.create_button("Clear", self.clear_start))
        form_layout.addRow("Start Point:", start_layout)

        # End Point
        end_layout = QtWidgets.QHBoxLayout()
        self.end_field = QtWidgets.QLineEdit()
        end_layout.addWidget(self.end_field)
        end_layout.addWidget(self.create_button("Assign", self.assign_end))
        end_layout.addWidget(self.create_button("Select", self.select_end))
        end_layout.addWidget(self.create_button("Clear", self.clear_end))
        form_layout.addRow("End Point:", end_layout)

        # Vertices
        vert_layout = QtWidgets.QHBoxLayout()
        self.vert_field = QtWidgets.QLineEdit()
        vert_layout.addWidget(self.vert_field)
        vert_layout.addWidget(self.create_button("Assign", self.assign_verts))
        vert_layout.addWidget(self.create_button("Select", self.select_verts))
        vert_layout.addWidget(self.create_button("Clear", self.clear_verts))
        form_layout.addRow("Vertices:", vert_layout)

        layout.addLayout(form_layout)

        # Order Radio Buttons
        order_layout = QtWidgets.QHBoxLayout()
        self.order_group = QtWidgets.QButtonGroup()
        self.passed_order = QtWidgets.QRadioButton("Use Passed Order")
        self.spatial_order = QtWidgets.QRadioButton("Use Spatial Order")
        self.order_group.addButton(self.passed_order)
        self.order_group.addButton(self.spatial_order)
        self.passed_order.setChecked(True)
        order_layout.addWidget(self.passed_order)
        order_layout.addWidget(self.spatial_order)
        layout.addLayout(order_layout)

        # Include Start/End Checkboxes
        include_layout = QtWidgets.QHBoxLayout()
        self.include_start_checkbox = QtWidgets.QCheckBox("Include Start")
        self.include_end_checkbox = QtWidgets.QCheckBox("Include End")
        self.include_start_checkbox.setChecked(True)
        self.include_end_checkbox.setChecked(True)
        include_layout.addWidget(self.include_start_checkbox)
        include_layout.addWidget(self.include_end_checkbox)
        layout.addLayout(include_layout)

        '''
        # Slope Editor
        self.slope_editor = self.create_slope_editor()
        layout.addWidget(QtWidgets.QLabel("Distribution Curve:"))
        layout.addWidget(self.slope_editor)

        # Curve Type Dropdown
        self.curve_type_combo = QtWidgets.QComboBox()
        self.curve_type_combo.addItems(["Linear", "Smooth", "Spline", "Exponential Up", "Exponential Down"])
        self.curve_type_combo.currentIndexChanged.connect(self.update_curve_type)
        layout.addWidget(self.curve_type_combo)
        '''

        # Align and Distribute Button
        self.align_button = QtWidgets.QPushButton("Align and Distribute")
        self.align_button.clicked.connect(self.align_and_distribute)
        self.align_button.setEnabled(False)
        layout.addWidget(self.align_button)

        self.setLayout(layout)

    def create_button(self, text, callback):
        button = QtWidgets.QPushButton(text)
        button.clicked.connect(callback)
        return button
    '''
    def create_slope_editor(self):
        # Create the gradientControlNoAttr
        gradient_control = cmds.gradientControlNoAttr(
            h=100,  # Height of the control
            w=400   # Width of the control
        )

        # Set initial keyframes for start and end
        cmds.gradientControlNoAttr(gradient_control, e=True, vap=0.0)
        cmds.gradientControlNoAttr(gradient_control, e=True, vap=1.0)

        # Convert the Maya UI element to a QWidget
        ptr = omui.MQtUtil.findControl(gradient_control)
        if ptr is not None:
            return wrapInstance(int(ptr), QtWidgets.QWidget)
        return None

    def update_curve_type(self, index):
        curve_type = self.curve_type_combo.currentText()
        if curve_type == "Linear":
            self.set_curve_interp(0)  # 0 for linear
        elif curve_type == "Smooth":
            self.set_curve_interp(1)  # 1 for smooth
        elif curve_type == "Spline":
            self.set_curve_interp(2)  # 2 for spline
        elif curve_type == "Exponential Up":
            self.set_curve_interp(3)  # 3 for exponential up
        elif curve_type == "Exponential Down":
            self.set_curve_interp(4)  # 4 for exponential down

    def set_curve_interp(self, interp_value):
        # Set the interpolation for all keys
        keys = cmds.gradientControlNoAttr(self.slope_editor, q=True, asString=True).split(",")
        num_keys = len(keys) // 4  # Each key is represented by 4 values
        for i in range(num_keys):
            cmds.gradientControlNoAttr(self.slope_editor, e=True, ckv=i, civ=interp_value)
    '''
    def clear_start(self):
        self.start_field.clear()
        self.update_align_button()

    def assign_start(self):
        sel = cmds.ls(sl=True, fl=True)
        if sel:
            self.start_field.setText(sel[0])
        self.update_align_button()

    def select_start(self):
        obj = self.start_field.text()
        if obj:
            cmds.select(obj, r=True)

    def clear_end(self):
        self.end_field.clear()
        self.update_align_button()

    def assign_end(self):
        sel = cmds.ls(sl=True, fl=True)
        if sel:
            self.end_field.setText(sel[0])
        self.update_align_button()

    def select_end(self):
        obj = self.end_field.text()
        if obj:
            cmds.select(obj, r=True)

    def clear_verts(self):
        self.vert_field.clear()
        self.update_align_button()

    def assign_verts(self):
        sel = cmds.ls(sl=True, fl=True)
        verts = cmds.filterExpand(sel, sm=31)
        if verts:
            self.vert_field.setText(','.join(verts))
        self.update_align_button()

    def select_verts(self):
        verts = self.vert_field.text().split(',')
        if verts:
            cmds.select(verts, r=True)

    def update_align_button(self):
        self.align_button.setEnabled(
            bool(self.start_field.text()) and
            bool(self.end_field.text()) and
            bool(self.vert_field.text())
        )

    def align_and_distribute(self):
        start = self.start_field.text()
        end = self.end_field.text()
        verts = self.vert_field.text().split(',')
        use_spatial = self.spatial_order.isChecked()
        include_start = self.include_start_checkbox.isChecked()
        include_end = self.include_end_checkbox.isChecked()

        # Check if start and end are the same object
        if start == end:
            cmds.warning("Start and end points are the same object. Please choose different objects.")
            return

        self.distribute_vertices(verts, start, end, use_spatial, include_start, include_end)

    def get_position(self, obj):
        if '.vtx[' in obj:
            vertex_index = int(obj.split('[')[-1].split(']')[0])
            mesh_name = obj.split('.')[0]

            sel_list = om.MSelectionList()
            sel_list.add(mesh_name)
            dag_path = om.MDagPath()
            sel_list.getDagPath(0, dag_path)

            fn_mesh = om.MFnMesh(dag_path)
            point = om.MPoint()
            fn_mesh.getPoint(vertex_index, point, om.MSpace.kWorld)

            return [point.x, point.y, point.z]
        else:
            return cmds.xform(obj, q=True, ws=True, t=True)

    def distribute_vertices(self, vertices, start, end, spatial=False, include_start=True, include_end=True):
        start_pos = self.get_position(start)
        end_pos = self.get_position(end)

        # Check if start and end positions are too close
        tolerance = 1e-5
        if self.distance(start_pos, end_pos) < tolerance:
            cmds.warning("Start and end points are too close. Please choose points further apart.")
            return

        if spatial:
            vertices = sorted(vertices, key=lambda v: self.distance(self.get_position(v), start_pos))

        total_distance = self.distance(start_pos, end_pos)

        # Calculate the number of segments based on include_start and include_end
        num_segments = len(vertices) - 1
        if not include_start:
            num_segments += 1
        if not include_end:
            num_segments += 1

        step = total_distance / num_segments

        # Start an undo chunk
        cmds.undoInfo(openChunk=True)
        try:
            for i, vertex in enumerate(vertices):
                # Calculate the t value based on include_start and include_end
                if include_start and include_end:
                    t = i * step / total_distance
                elif include_start and not include_end:
                    t = i * step / total_distance
                elif not include_start and include_end:
                    t = (i + 1) * step / total_distance
                else:  # not include_start and not include_end
                    t = (i + 1) * step / total_distance

                new_pos = [
                    start_pos[j] + t * (end_pos[j] - start_pos[j])
                    for j in range(3)
                ]
                cmds.xform(vertex, ws=True, t=new_pos)
        finally:
            # Close the undo chunk
            cmds.undoInfo(closeChunk=True)

    def distance(self, p1, p2):
        return sum((a - b) ** 2 for a, b in zip(p1, p2)) ** 0.5


def show_vertex_distributor():
    global vertex_distributor_window
    try:
        vertex_distributor_window.close()
        vertex_distributor_window.deleteLater()
    except:
        pass
    vertex_distributor_window = VertexDistributor()
    vertex_distributor_window.show()

show_vertex_distributor()