import prman
import math
  
class VertexDB(object):
    def __init__(self, nloops, nverts, verts, params):
        self.nloops = nloops
        self.nverts = nverts
        self.verts = verts
        self.xyz_verts = params[prman.Ri.P]
        self.xyz_norms = params['facevarying normal N']
        # This dictionary uses vertices as keys. The values are a single 
        # tuple that represents the average normal of the polygon faces that 
        # "share" a vertex (key).
        self.__averagedNormalsDB = {}
        self.__getDB()
        
        #self.__getAveNormalDB()
    #------------------------------------------------------------
    def __getDB(self):
        # Each key in the dictionary is a tuple of x,y,z (vertex) coordinates.
        # The values are a list of tuples of x,y,z (normal) coordinates of 
        # polygon faces that "share" vertex (key).
        sharedVertices = {}
        loop_counter = 0
        vert_counter = 0
        norm_index = 0     
        for n in range(len(self.nloops)):
            for j in range(self.nloops[n]):
                num_verts = self.nverts[loop_counter]
                loop_counter += 1
                for k in range(num_verts):
                    vert_index = self.verts[vert_counter] * 3
                    vert_x = self.xyz_verts[vert_index + 0]
                    vert_y = self.xyz_verts[vert_index + 1]
                    vert_z = self.xyz_verts[vert_index + 2]
                
                    norm_x = self.xyz_norms[norm_index + 0]
                    norm_y = self.xyz_norms[norm_index + 1]
                    norm_z = self.xyz_norms[norm_index + 2]
                    norm_index += 3
                    xyz = (vert_x, vert_y, vert_z)
                    if sharedVertices.has_key(xyz):
                        sharedVertices[xyz].append( (norm_x, norm_y, norm_z) )
                    else:
                        sharedVertices[xyz] = [ (norm_x, norm_y, norm_z) ]        
                    vert_counter += 1
        for key in sharedVertices.keys():
            X = Y = Z = 0.0
            normals = sharedVertices[key]
            num = len(normals)
            for norm in normals:
                X += norm[0]
                Y += norm[1]
                Z += norm[2]
            self.__averagedNormalsDB[key] = (X/num, Y/num, Z/num)
    #------------------------------------------------------------
    def getVertices(self):
        return self.__averagedNormalsDB.keys()
        
    def getAveNormal(self, vertex):
        return self.__averagedNormalsDB[vertex]    
    
    #------------------------------------------------------------
    # Based on the mel proc aimY() explained at
    #    http://fundza.com/mel/axis_to_vector/index.html
    def getAlignmentRotations(self, vertex):
        nx,ny,nz = self.__averagedNormalsDB[vertex]
        transforms = []
        xy_length = math.sqrt(nx * nx + ny * ny)
        norm_length = math.sqrt(nx * nx + ny * ny + nz * nz)
        if xy_length == 0:
            if nx > 0:
                z_angle = math.radians(90.0)
            else:
                z_angle = math.radians(-90.0)
        else:
            z_angle = math.acos( ny/xy_length )
  
        if nz > 0:
            x_angle = math.acos( xy_length/norm_length )
        else:
            x_angle = -math.acos( xy_length/norm_length )
        transforms.append( math.degrees(x_angle) )
        if nx > 0:
            z_angle *= -1.0
        transforms.append( math.degrees(z_angle) )
        return transforms