#!/usr/bin/env python
# Author:  Webber Huang
# E-mail:  xracz.fx@gmail.com
# Website: http://riggingtd.com
# Purpose: Find the closest edge on mesh from any position
# Created: 2014/6/17
#========================================================================

import heapq
import maya.cmds as cmds
import maya.OpenMaya as om

########################################################################
class DistBaseVertex(object):
    #----------------------------------------------------------------------
    def __init__(self, idx, dist):
        self._dist = dist
        self._idx = idx
        self._edges = set()
    
    #----------------------------------------------------------------------
    def index(self):
        return self._idx
    
    #----------------------------------------------------------------------
    def addConnectedEdges(self, idx):
        self._edges.add(idx)
    
    #----------------------------------------------------------------------
    def getConnectedEdges(self):
        return self._edges
    
    #----------------------------------------------------------------------
    def __repr__(self):
        return ("Index: %d\tDistance: %f\tConnected Edges: %s"
                % (self._idx, self._dist, str(self._edges)[5:-2]))
    
    #----------------------------------------------------------------------
    def __cmp__(self, other):
        """"""
        assert isinstance(other, DistBaseVertex)
        return cmp(self._dist, other._dist)


########################################################################
class PriorityQueue(object):
    #----------------------------------------------------------------------
    def __init__(self):
        self._queue = []
        self._count = 0
    
    #----------------------------------------------------------------------
    def push(self, item):
        heapq.heappush(self._queue, item)
        self._count += 1
    
    #----------------------------------------------------------------------
    def pop(self):
        result = heapq.heappop(self._queue)
        self._count -= 1
        return result
    
    #----------------------------------------------------------------------
    def count(self):
        return self._count
    
    #----------------------------------------------------------------------
    def __repr__(self):
        temp = list(self._queue)
        result = ""
        while temp: result += "%s\n" % heapq.heappop(temp)
        return result

#----------------------------------------------------------------------    
def toDagPathObj(name):
    sels = om.MSelectionList()
    sels.add (name)
    dagPath = om.MDagPath()
    sels.getDagPath(0, dagPath)
    return dagPath

#----------------------------------------------------------------------
def closestFaceVerticesOnMesh(mesh, position):
    meshObj = toDagPathObj(mesh)    
    meshFn = om.MFnMesh(meshObj)
    
    srcPt = om.MPoint(position[0], position[1], position[2])
    tgtPt = om.MPoint()
    
    util = om.MScriptUtil()
    idPtr = util.asIntPtr()
    
    meshFn.getClosestPoint(srcPt, tgtPt, om.MSpace.kWorld, idPtr)
    idx = om.MScriptUtil.getInt(idPtr)

    vertices = om.MIntArray()
    meshFn.getPolygonVertices(idx, vertices)
    
    return vertices

#----------------------------------------------------------------------
def closestEdgeOnMesh(mesh, position, returnId=False):
    meshObj = toDagPathObj(mesh)    
    vtxIt = om.MItMeshVertex(meshObj)
    util = om.MScriptUtil()
    prevIdx = util.asIntPtr()     
    
    srcPt = om.MPoint(position[0], position[1], position[2])
    pq = PriorityQueue()
    
    # Get closest face vertices
    vertices = closestFaceVerticesOnMesh(mesh, position)
    
    # Iterate vertices, create DistBaseVertex object, push into pq
    for v in vertices:
        vtxIt.reset()
        vtxIt.setIndex(v, prevIdx)
        
        # Measure distance from source point to current vertex
        dist = vtxIt.position(om.MSpace.kWorld).distanceTo(srcPt)
        dbv = DistBaseVertex(v, dist)
        
        # Add connected edges
        edges = om.MIntArray()
        vtxIt.getConnectedEdges(edges)
        for e in edges: dbv.addConnectedEdges(e)
        
        pq.push(dbv)
    
    # Obtain closest edge by calculate the intersection of two edge sets
    vtx1, vtx2= pq.pop(), pq.pop()
    idx = (vtx1.getConnectedEdges() & vtx2.getConnectedEdges()).pop()
    
    if returnId: return idx
    return "%s.e[%d]" % (mesh, idx)
    
if __name__ == "__main__":
    src, mesh = cmds.ls(sl=1)
    pos = cmds.xform(src, q=1, ws=1, t=1)
    e = closestEdgeOnMesh(mesh, pos, False)
    cmds.select(e)
    om.MGlobal.displayInfo(e)