#!BPY

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

__author__  = "Guillaume Englert"
__version__ = "0.3 2009/07/15"
__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.<br>
 - Rectify : rotate UV groups in order to have vertical and horizontal edges.
"""

################################################################################
#                                                                              #
#    GNU GPL LICENSE                                                           #
#    ---------------                                                           #
#                                                                              #
#    Copyright (C) 2006-2009: 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


################################################################################
#                              CLASSES                                         #
################################################################################

class UVToolAction(object):
    """Action that can do this module, like projection or intersection."""
    def __init__(self, function, label, **kw):
        """Constructor.
        function: function object that do the action
        label: description of the action (for menu)
        """
        self.__function = function
        self.__label    = label

        for key, val in kw.iteritems():
            setattr(self, key, val)

    def __call__(self, *args):
        return self.__function(self, *args)

    @property
    def label(self):
        return self.__label

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

class FaceUV(object):
    """Owner of an UV coordinate."""
    __slots__ = ('faceid', 'uvid')
    def __init__(self, faceid, uvid):
        """Constructor.
        faceid: index of the face.
        uvid: index of the UV coord for this face.
        """
        self.faceid = faceid
        self.uvid   = uvid

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

class UVPoint(object):
    """An UV coordinate, which is shared by some faces."""
    def __init__(self, uv, faceuv):
        """Constructor.
        uv: UV coordinates (Blender.Mathutils.Vector).
        faceuv: a FaceUV object which corresponds to this UV point.
        """
        self.uv      = uv
        self.facelst = [faceuv]  #list of FaceUV objects which share this UV coordinate


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

EPSILON = 0.001


################################################################################
#                             GLOBALS                                          #
################################################################################

actions = [] #list of UVToolAction objects


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

##############
# DECORATORS #
##############

def actionfunc(label, **kw):
    """Decorator that builds a UVToolAction which have a method that corresponds
    to the decoated function, and adds the new action to the global actions list.
    label: label that describe the function (used in the popup menu).
    kw: keyword parameter will be added as members to the UVToolAction object.
    Notice that decorated function are on the model:
    def foobar(self, mesh):
        ...
    where self is the UVToolAction object, and mesh the Blender.Mesh.Mesh concerned.
    """
    def _deco(func):
        actions.append(UVToolAction(func, label, **kw))
        return func
    return _deco


##################
# USER INTERFACE #
##################

def get_action():
    """Ask to the user the action he wants.
    return: None or a UVToolAction object
    """
    ret = PupMenu('UV tool%t|' + '|'.join(action.label for action in actions))

    if ret == -1: #Cancel
        return None
    else:
        return actions[ret-1]

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

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


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


#########
# UTILS #
#########

def build_uvpoints(mfaces, uvptclass=UVPoint):
    """Build the list of UVPoint objects.
    mfaces: faces list of the mesh.
    uvptclass: class of UV point to use (UVPoint or a child class)
    return: a list of uvptclass objects.
    """
    uvpts = [] #UV Points

    for face in mfaces:
        if not face.sel:
            continue

        for i, flag in enumerate(face.uvSel):
            if not flag: continue

            uv = face.uv[i]

            for uvpt in uvpts:
                if (uv - uvpt.uv).length < EPSILON:
                    uvpt.facelst.append(FaceUV(face.index, i))
                    break
            else:
                uvpts.append(uvptclass(uv, FaceUV(face.index, i)))

    return uvpts

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


#######
# UV  #
#######


@actionfunc("Super weld")
def superweld(self, 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()

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

@actionfunc("Distribute")
def distribute(self, mesh):
    """Distribute regularly UV coordinates beetween the 2 extremities.
    mesh: selected mesh (Blender.Mesh.Mesh object).
    """
    mfaces = mesh.faces
    uvpts  = build_uvpoints(mfaces)

    if len(uvpts) < 3: return

    #sort the UV points 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
    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()

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

@actionfunc("Rectify")
def rectify(self, mesh):
    from Blender.Mathutils import Vector, AngleBetweenVecs, RotationMatrix, Matrix
    from itertools import islice

    class _UVPoint(UVPoint):
        def __init__(self, uv, faceuv):
            UVPoint.__init__(self, uv, faceuv)
            self.faceset = None     #set of the faces indices

        def compute_faceset(self):
            self.faceset = set(faceuv.faceid for faceuv in self.facelst)

    #---------------------------------------------------------------------------
    class _Forest:
        """A group of connexity trees, in order to know the connex group"""
        def __init__(self, links_lst):
            """Constructor.
            links_lst: list of links ((int, int) tuples).
            """
            # key: value   --> value is the parent of key (for roots, key is None).
            self._links = dict((v, None) for link in links_lst for v in link)
            #key: value --> value is the size of the tree whose root is key.
            self._sizes = dict((key, 1) for key in self._links.iterkeys())

            for v1, v2 in links:
                self._union(v1, v2)

        def _root(self, key):
            links  = self._links
            tmpkey = links[key]
            while tmpkey is not None:
                key    = tmpkey
                tmpkey = links[key]

            return key

        def _union(self, key1, key2):
            root1 = self._root(key1)
            root2 = self._root(key2)

            if root1 != root2:
                size1 = self._sizes[root1]
                size2 = self._sizes[root2]
                if size1 < size2:
                    self._links[root1]  = root2
                    self._sizes[root2] += size1
                else:
                    self._links[root2]  = root1
                    self._sizes[root1] += size2

        def get_groups(self):
            """Get the connexity groups.
            return: a list of groups. A group is a list of ints.
            """
            trees = dict((k, []) for k, v in self._links.iteritems() if v is None)

            for k in self._links.iterkeys():
                trees[self._root(k)].append(k)

            return trees.values()

    #---------------------------------------------------------------------------
    def _build_links(uvpts):
        """Build the links between UV points.
        uvpts: list of UVPoint objects.
        return: list of links. A link is a (int, int) tuple, where integers are
        indices of UV points in uvpts."""
        def _find_uvid(uvpt, faceid):
            for uvface in uvpt.facelst:
                if faceid == uvface.faceid:
                    return uvface.uvid

        def _link_gen():
            for i, uvpt1 in enumerate(islice(uvpts, len(uvpts)-1)):
                faceset1 = uvpt1.faceset

                for j, uvpt2 in enumerate(islice(uvpts, i+1, len(uvpts))):
                    for faceid in faceset1 & uvpt2.faceset:
                        diff = abs(_find_uvid(uvpt1, faceid) - _find_uvid(uvpt2, faceid))

                        if (diff == 1) or (diff == len(mfaces[faceid].uv)-1):
                            yield (i, j+i+1)
                            break

        for uvpt in uvpts:
            uvpt.compute_faceset()

        return [link for link in _link_gen()]

    #---------------------------------------------------------------------------
    def _angle(vect):
        """Angle of a vector ; angle is beetween -45. and 45 degrees.
        vect: vector (Blender.Mathutils.Vector object).
        return: the angle.
        """
        if vect.x < 0.:
            vect = -vect

        angle = AngleBetweenVecs(unitv, vect)

        if angle >= 45.: #-45 =< angle < 45
            angle -= 90.

        if vect.y >= 0.:
            return -angle

        return angle


    def _compute_bestangle(vects):
        """Get the best angle with a list of vector.
        vects: a list of Blender.Mathutils.Vector objects.
        return: the best angle
        """
        angles = [_angle(vect) for vect in vects]
        angles.sort()

        bestangle = None #0.0

        maxcount = 1
        curcount = 1      #current count
        epsilon  = 0.001

        it       = iter(angles)
        oldangle = it.next()

        for angle in it:
            if (angle-oldangle) < epsilon:
                curcount += 1
                if curcount > maxcount:
                    maxcount  = curcount
                    bestangle = angle
            else:
                curcount = 1
            oldangle = angle

        if bestangle is None: #take the longest vector
            maxlen = -1.0
            for vect in vects:
                vlen = vect.length
                if vlen > maxlen:
                    maxlen   = vlen
                    bestvect = vect

            bestangle = _angle(bestvect)

        return bestangle

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

    mfaces     = mesh.faces
    uvpts      = build_uvpoints(mfaces, _UVPoint)
    links      = _build_links(uvpts)
    clusters   = _Forest(links).get_groups()
    clusterdic = dict((i, j) for j, cluster in enumerate(clusters) for i in cluster)

    vectlst = [[] for i in clusters]
    for l1, l2 in links:
        vectlst[clusterdic[l1]].append(uvpts[l2].uv - uvpts[l1].uv)

    nullv  = Vector(0.0, 0.0) #null vector
    unitv  = Vector(1.0, 0.0) #x axis unit vector

    for clustind, cluster in enumerate(clusters):
        center    = sum((uvpts[i].uv for i in cluster), nullv) * (1.0/len(cluster))
        bestangle = _compute_bestangle(vectlst[clustind])

        mat  = Matrix([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-center.x, -center.y, 1.0]) #translation matrix
        mat *= RotationMatrix(bestangle, 3, 'z')
        mat *= Matrix([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [center.x, center.y, 1.0])   #translation matrix

        for i in cluster:
            uvpt     = uvpts[i]
            uvvect   = uvpt.uv
            finaluv  = Vector(uvvect.x, uvvect.y, 1.0)
            finaluv *= mat
            finaluv.resize2D()

            for face in uvpt.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:
            action(mesh)


    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()