UV Unwrap KD-Tree, NNS, A* Vertex/Edge Selection

'''
    WIP: 15_1206
    Script By: Dr. Gravitas

    Imports a pickled representation of edges and their vertices, representing the
    borders of UV Shells. This tries to replicate the actions taken to create a
    UV based on the nearest edge/vertex of the selected model. It sews together
    all of the UVs and then cuts those nearest edges and finally unwraps it.

'''
import pymel.core as pm
from maya import OpenMaya as om
import os
import fileinput
import pickle
import heapq
from collections import namedtuple
from operator import itemgetter
from pprint import pformat


'''
    Given a vector and a KDTree Node, look up the vertex corresponding with the node's location and
    determine the distance in worldspace between it and the vector's point.
'''
def compareVectorDistanceFromNode(vector, node):
    if node is not None:
        vertex = v2PDict[node.location]
        return vector.distanceTo( pm.datatypes.Vector( vertex.getPosition(space='world') ) )
    # When the node is none, return infinite distance
    return float('inf')


def compareVectorDistance(vertex, vector):
    return vector.distanceTo( pm.datatypes.Vector( vertex.getPosition(space='world') ) )

class PathNode:
    cost, node, edge, parentNode = (None,None,None,None)

    def __init__(self, cost=None,node=None,edge=None,parentNode=None):
        self.cost = cost
        self.node = node
        self.edge = edge
        self.parentNode = parentNode

def getConnectedVertsAndEdges( vertex ):
    connectedPathNodes = []
    edges = vertex.connectedEdges()

    for edge in edges:
        verts = edge.connectedVertices()
        vert = verts[(verts.index(vertex) -1)*-1]
        #Uninitialized cost is infinite, parent is none
        connectedPathNode = PathNode( node=vert, edge=edge )
        connectedPathNodes.append( connectedPathNode )

    return connectedPathNodes

def retracePath( pathNode ):
    path = [pathNode.edge]
    while pathNode.parentNode is not None:
        pathNode = pathNode.parentNode
        path.append(pathNode.edge)
    path.reverse()
    return path

def aStar( startPathNode, end, expandNodeMethod, comparisonMethod ):
    openSet = set()
    openHeap = []
    closedSet = set()

    openSet.add(startPathNode.node)
    openHeap.append(startPathNode)

    while openSet:
        currentPathNode = heapq.heappop(openHeap)
        if comparisonMethod(currentPathNode.node, end) == 0:
            return retracePath(currentPathNode)

        openSet.remove(currentPathNode.node)
        closedSet.add(currentPathNode.node)

        for connectedNode in expandNodeMethod( currentPathNode.node ):
            if connectedNode.node not in closedSet:
                # Estimated distance to the end
                dist = comparisonMethod( connectedNode.node, end )
                if connectedNode.node not in openSet:
                    openSet.add(connectedNode.node)
                    connectedNode.cost = dist
                    heapq.heappush(openHeap, connectedNode)
                connectedNode.parentNode = currentPathNode
    return None

'''
    Get minimum edges necessary to connect the two vertices
'''
def pathBetweenVertices( startVertex, goalVertex ):
    edges = []

    if startVertex.index() == goalVertex.index():
        pm.error('Identical Vertices provided for pathfinding!' + str(startVertex) + ' and ' + str(goalVertex) )
    if not startVertex.isConnectedTo( goalVertex ):
        #For when the two vertices are not directly connected
        startNode = PathNode( cost = 0,  node = startVertex)
        goalAsVector = pm.datatypes.Vector( goalVertex.getPosition(space='world') )
        edges = aStar( startNode, goalAsVector, getConnectedVertsAndEdges, compareVectorDistance )
        if edges is None:
            pm.warning('A* failed to connect the two vertices: ' + str(startVertex) + ' and ' + str(goalVertex) )
            edges = []
    else:
        possibleEdges = startVertex.connectedEdges()
        for possibleEdge in possibleEdges:
            if goalVertex in possibleEdge.connectedVertices():
                edges.extend( possibleEdge )

    return edges



class BestChoice(namedtuple('BestChoice', 'node distanceToIdeal')):
    def __repr__(self):
        return pformat(tuple(self))

def recursive_search( idealPoint, illegalPoints, currentNode, bestNode, comparisonMethod ):
    if currentNode is None:
        return bestNode

    # Do not allow the vertices that are listed to be selected as the best point.
    if currentNode.location not in illegalPoints:
        distance = comparisonMethod( idealPoint, currentNode )
        if distance < bestNode.distanceToIdeal:
            #print distance, bestNode.distanceToIdeal, currentNode.location
            bestNode = BestChoice( currentNode, distance )

    # Determine axis (X, Y, Z) of the current node (depth mod number of dimensions)
    axis = currentNode.axis
    # Get the difference between the idealPoint and the current Node's position at the chosen axis
    #  and determine whether to travel down the left or right branch
    diff = idealPoint[axis] - currentNode.location[axis]

    close, away = (currentNode.left_child, currentNode.right_child) if diff <= 0 else (currentNode.right_child, currentNode.left_child)

    bestNode = recursive_search( idealPoint, illegalPoints, close, bestNode, comparisonMethod )
    if diff ** 2 < bestNode.distanceToIdeal:
        bestNode = recursive_search(idealPoint, illegalPoints, away, bestNode, comparisonMethod )

    return bestNode


def searchKDTree( idealPoint, illegalPoints, comparisonMethod ):
    #Execute the search at the root node of the tree
    distance = comparisonMethod( idealPoint, tree )
    bestNode = BestChoice( tree, distance )
    bestNode = recursive_search( idealPoint, illegalPoints, tree, bestNode, comparisonMethod )
    return bestNode.node

'''
    Worst Case: We have no idea where this edge begins or ends.

    Given an edgeDef, find the edge or edges that best fits that definition.
'''
def searchForEdge(edgeDef):
    actualEdges = []
    verts = edgeDef[1]
    actualVerts = []
    illegalPoints = []

    for vertDef in verts:
        vertPoint = vertDef[1]
        vector = pm.datatypes.Vector( vertPoint )
        # Begin KDTree navigation - identified vertices are illegal to reselect for subsequent vertices of the edge
        bestNode = searchKDTree( vertPoint, illegalPoints, compareVectorDistanceFromNode )
        #print bestNode.tree_level
        bestVert = v2PDict[bestNode.location]
        actualVerts.extend( bestVert )
        illegalPoints.append( bestNode.location )
        #print vertDef
        #print '[' + str(bestVert) + ', ' + str(bestVert.getPosition())

    #With best cases for Verts, get edges between first and all others. Should only be one other!
    start = actualVerts[0]
    for end in actualVerts[1:]:
        actualEdges.extend( pathBetweenVertices( start, end ) )


    return actualEdges

'''
    Identifies MeshEdges for each of the edges definied in the edgeDefList
    and cuts the UV along them.
'''
def getActualEdges(edges):
    actualEdges = []

    for edge in edges:
        #actualEdge = findActualEdge( edge )
        #actualEdges.extend( actualEdge )
        actualEdges.extend( searchForEdge( edge ) )

    return actualEdges

'''
    Attempts to reconstruct UVs based on list of edges/vertices
'''
def reconstructUVOps(selected):
    pm.polyMapSew( selected )
    edges = importUVEdgeDefs()
    actualEdges = getActualEdges( edges )
    pm.select( clear=True )
    pm.polyMapCut( actualEdges )
    pm.select( target.map[0:] )
    pm.mel.eval('Unfold3D -u -ite 3 -p 1 -bi 1 -tf 1 -ms 4096 -rs 3;')
    pm.select( actualEdges )
    pm.refresh(f=True)

'''
    Attempts to read in from a file and unpickle the data into a Python data structure
'''
def readFile(inputFile):
    commands = None
    try:
        with open(inputFile, 'r') as inFile:
            commands = pickle.load( inFile )
    except IOError:
        pm.error("Error: File does not appear to exist.")
    return commands


def buildTreeAndHash( target ):
    for vertex in target.verts:
        vPoint = vertex.getPosition( space='world' )
        pos = ( vPoint.x, vPoint.y, vPoint.z )

        v2PDict[pos] = vertex

    tree = kdtree( v2PDict.keys() )
    return tree

class Node(namedtuple('Node', 'location tree_level axis left_child right_child')):
    def __repr__(self):
        return pformat(tuple(self))

def kdtree( point_list, depth=0 ):
    try:
        k = len( point_list[0] ) # assumes all points have the same dimension
    except IndexError as e: # if not point_list:
        return None
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list and choose median as pivot element
    point_list = sorted( point_list, key=itemgetter(axis))
    median = len( point_list ) // 2 # choose median

    # Create node and construct subtrees
    return Node(
        location = point_list[median],
        tree_level = depth,
        axis = axis,
        left_child = kdtree(point_list[:median], depth + 1),
        right_child = kdtree(point_list[median + 1:], depth + 1)
    )

'''
    Imports Edges from a file location into Python Data structures

    TODO: Turn this into a dialog box that lets you configure this but auto-populates
    with most recent file in default location

'''
def importUVEdgeDefs():
    dirPath = os.path.dirname( 'C:\\[Insert Path to Project Scripts Here]\\Scripts\\StoredData\\UV Cut Edges\\' )
    filename = '15_1206_UVCutEdges' # Name of the exported edges definition
    inputFile = os.path.abspath( os.path.join( dirPath, filename ) )

    return readFile( inputFile )


'''
    Main Body
'''
start_time = pm.date()
print '\n\n\n'
print '******************************************************************'
print '\t\t\t' + start_time
print '******************************************************************'
target = (pm.selected())[0].getShape()
#pm.select( clear=True )
tree = ()
v2PDict = {}
tree = buildTreeAndHash( target )
reconstructUVOps( target )
tree = None
v2PDict = None
end_time = pm.date()
print '******************************************************************'
print '\t\t\t' + start_time
print '\t\t\t' + end_time
print '******************************************************************'
#pm.error('Success!')