#!BPY

"""
Name: 'UV Tool'
Blender: 248
Group: 'UV'
Tooltip: 'Some tools to manage UVs (super weld, distribute)'
"""

__author__  = "Guillaume Englert"
__version__ = "0.1 2006/10/20"
__url__     = "Online doc , http://www.hybird.org/~guiea_7/"
__email__   = "GuieA_7, genglert:hybird*org"
__bpydoc__  = """\
Some tools to manage UVs:<br>
 - Super weld : weld the selected UV by group, like 'remove doubles'.<br>
 - Distribute : distribute regularly UV points which are (nearly) aligned.
"""

################################################################################
#                                                                              #
#    GNU GPL LICENSE                                                           #
#    ---------------                                                           #
#                                                                              #
#    Copyright (C) 2006: Guillaume Englert                                     #
#                                                                              #
#    This program is free software; you can redistribute it and/or modify it   #
#    under the terms of the GNU General Public License as published by the     #
#    Free Software Foundation; either version 2 of the License, or (at your    #
#    option) any later version.                                                #
#                                                                              #
#    This program is distributed in the hope that it will be useful,           #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of            #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the             #
#    GNU General Public License for more details.                              #
#                                                                              #
#    You should have received a copy of the GNU General Public License         #
#    along with this program; if not, write to the Free Software Foundation,   #
#    Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.           #
#                                                                              #
################################################################################


################################################################################
# Importing modules
################################################################################

from Blender import Mesh, Redraw
from Blender.Object import GetSelected
from Blender.Window import EditMode
from Blender.Draw import PupMenu

from exceptions import Exception
from itertools import islice


################################################################################
#                              EXCEPTION                                       #
################################################################################

class UVToolError(Exception):
    def __init__(self, msg):
        Exception.__init__(self)
        self.msg = msg

    def __str__(self):
        return self.msg


################################################################################
#                             CONSTANTS                                        #
################################################################################

EPSILON = 0.001

SUPERWELD  = 1
DISTRIBUTE = 2
CANCEL     = -1

################################################################################
#                             FUNCTIONS                                        #
################################################################################


def display_error(string):
    PupMenu("ERROR: " + string)

#-------------------------------------------------------------------------------

def get_action():
    return PupMenu(
    "UV tools%t|"
    "Super weld|"
    "Distribute")

#-------------------------------------------------------------------------------

def superweld(mesh):
    """Find groups of UV coordinates that are close, and weld them.
    mesh: selected mesh (Blender.Mesh.Mesh object).
    """
    from Blender.Mathutils import Vector
    from Blender.Draw import PupFloatInput

    #get the limit for the weld
    epsilon = PupFloatInput("Limit", 0.001, 0.0001, 1.0, 0.1, 4)
    if not epsilon: return


    class _TaggedUV(object):
        __slots__ = ('uv', 'flag', 'faceid', 'uvid')
        def __init__(self, uv, faceid, uvid):
            self.uv     = uv     #UV coord (Vector)
            self.flag   = False  #if the object belongs to a group
            self.faceid = faceid #index of the face
            self.uvid   = uvid   #index of the UV coord for this face


    uvlist = [_TaggedUV(face.uv[i], face.index, i)
                for face in mesh.faces
                    for i, flag in enumerate(face.uvSel) if flag]
    uvlist.sort(key=lambda v: v.uv.x) #X coord sort

    groups = []
    for i1, uv1 in enumerate(islice(uvlist, len(uvlist)-1)):
        if uv1.flag: continue

        group = [uv1]
        x     = uv1.uv.x

        for uv2 in islice(uvlist, i1, len(uvlist)):
            if abs(x - uv2.uv.x) > epsilon: break
            if uv2.flag: continue

            if (uv1.uv - uv2.uv).length < epsilon:
                #if UVs are close --> group them
                uv2.flag = True
                group.append(uv2)

        groups.append(group)

    mfaces = mesh.faces
    nullv  = Vector(0.0, 0.0)
    for group in groups:
        #final UV = average of UV coordinates of the group
        finaluv = sum((tuv.uv for tuv in group), nullv) * (1.0/len(group))

        for tuv in group:
            uv   = mfaces[tuv.faceid].uv[tuv.uvid]
            uv.x = finaluv.x
            uv.y = finaluv.y

    mesh.update()

#-------------------------------------------------------------------------------

def distribute(mesh):
    """Distribute regularly UV coordinates beetween the 2 extremities.
    mesh: selected mesh (Blender.Mesh.Mesh object).
    """
    class _FaceUV:
        """Owner of an UV coordinate"""
        def __init__(self, faceid, uvid):
            self.faceid = faceid #index of the face
            self.uvid   = uvid   #index of the UV coord for this face

    class _UVPoint:
        """An UV coordinate, which is shared by some faces."""
        def __init__(self, uv, faceuv):
            self.uv      = uv       #UV coord (Vector)
            self.facelst = [faceuv] #list of _FaceUV which share that UV coordinate

    #build the list of _UVPoint objects
    uvpts = [] #UV Points
    for face in mesh.faces:
        for i, flag in enumerate(face.uvSel):
            if not flag: continue

            uv    = face.uv[i]
            found = False
            for uvpt in uvpts:
                if (uv - uvpt.uv).length < EPSILON:
                    uvpt.facelst.append(_FaceUV(face.index, i))
                    found = True
                    break

            if not found:
                uvpts.append(_UVPoint(uv, _FaceUV(face.index, i)))

    if len(uvpts) < 3: return

    #sort the UV point to get a list of aligned point (normally :)
    uvpts.sort(key=lambda pt: pt.uv.x)    #X coord sort
    uvptstmp = list(uvpts)
    uvptstmp.sort(key=lambda pt: pt.uv.y) #Y coord sort

    #compare the x diff and the y diff
    if abs((uvpts[0].uv.x - uvpts[-1].uv.x)) < abs((uvptstmp[0].uv.y - uvptstmp[-1].uv.y)):
        uvpts = uvptstmp
    del uvptstmp

    #UV points are regularly aligned
    mfaces = mesh.faces
    point  = uvpts[0].uv
    vect   = (uvpts[-1].uv - point) * (1.0/(len(uvpts)-1))

    for mult, pt in enumerate(uvpts):
        finaluv = mult * vect + point

        for face in pt.facelst:
            uv   = mfaces[face.faceid].uv[face.uvid]
            uv.x = finaluv.x
            uv.y = finaluv.y

    mesh.update()


################################################################################
#                           MAIN FUNCTION                                      #
################################################################################

def main():
    #init
    is_editmode = EditMode()
    if is_editmode:
        EditMode(0)

    try:
        #get selected object (or quit)
        objs = GetSelected()
        if not objs:
            raise UVToolError("none selected object")

        if len(objs) > 1:
            raise UVToolError("only one object must be selected")

        obj = objs[0]
        if obj.getType() != "Mesh":
            raise UVToolError("active object must be a mesh")

        mesh = obj.getData(mesh=True)

        if not mesh.faceUV:
            raise UVToolError("active object must have UV")

        #do the action
        action = get_action()

        if action == CANCEL:
            pass
        elif action == SUPERWELD:
            superweld(mesh)
        elif action == DISTRIBUTE:
            distribute(mesh)
        else:
            raise UVToolError("bad menu item")


    except UVToolError, e:
        print e
        display_error(str(e))

    except:
        import sys
        sys.excepthook(*sys.exc_info())
        display_error("An exception occured | (look at the terminal)")

    #finish
    if is_editmode :
        EditMode(1)

################################################################################
#                           MAIN PROGRAM                                       #
################################################################################

main()
