import numpy as np
import torch

class universe():
    def __init__(self,elements):
        self.elements = elements
        self.num = elements
        self.rank = np.zeros((self.elements,),dtype= int)
        self.size = np.ones((self.elements,),dtype= int)
        self.p = np.arange(self.elements)
    def find(self,x):
        y = x
        while y!= self.p[y]:
            y = self.p[y]
        self.p[x] = y
        return y
    def join(self,x,y):
        if self.rank[x] > self.rank[y]:
            self.p[y] = x
            self.size[x] += self.size[y]
        else:
            self.p[x] = y
            self.size[y] += self.size[x]
            if self.rank[x] == self.rank[y]:
                self.rank[y] += 1
        self.num -= 1
    def get_size(self,x):
        return self.size[x]
    def num_sets(self):
        return self.num


class vec3f():
    def __init__(self,_x=0.,_y=0.,_z=0.):
        self.x = _x
        self.y = _y
        self.z = _z 
    def vec_plus(self,vec2):
        return vec3f(self.x+vec2.x,self.y+vec2.y,self.z+vec2.z)
    def vec_sub(self,vec2):
        return vec3f(self.x-vec2.x,self.y-vec2.y,self.z-vec2.z)

class edge():
    def __init__(self):
        self.w = 0.
        self.a = int(0)
        self.b = int(0)

def cross(u,v):
    c = vec3f(u.y * v.z - u.z * v.y, u.z * v.x - u.x * v.z, u.x * v.y - u.y * v.x)
    n = (c.x * c.x + c.y * c.y + c.z * c.z)**0.5
    c.x /= n
    c.y /= n
    c.z /= n
    return c
def lerp(a,b,v):
    u = 1.0 - v
    return vec3f(v * b.x + u * a.x, v * b.y + u * a.y, v * b.z + u * a.z) 

def segment_graph(num_vertices, num_edges, edges, c):
    edges.sort(key=lambda x:x.w, reverse=False) # sort edges by weight
    u = universe(num_vertices) # make a disjoint-set forest
    threshold = np.zeros((num_vertices,),dtype = float)
    for i in range(num_vertices):
        threshold[i] = c
    # for each edge, in non-decreasing weight order
    for i in range(num_edges):
        pedge = edges[i]
        # components conected by this edge
        a = u.find(pedge.a)
        b = u.find(pedge.b)
        if a != b:
            if (pedge.w <= threshold[a]) and (pedge.w <= threshold[b]):
                u.join(a,b)
                a = u.find(a)
                threshold[a] = pedge.w + ( c / u.get_size(a))
    del threshold
    return u            
    
def segment_mesh_kernel(verts,vertexCount,faces,faceCount,kthr,segMinVerts):

    # create points, normals, edges, counts vectors
    points = []
    normals = []
    counts = np.zeros((vertexCount*3,),dtype=int)
    edgeCount = faceCount * 3
    edges = []
    for i in range(vertexCount):
        points.append(vec3f())
        normals.append(vec3f())
    for i in range(edgeCount):
        edges.append(edge())    

    

    # Compute face normals and smooth into vertex normals
    for i in range(faceCount):

        i1,i2,i3 = faces[i]
        p1 = vec3f(verts[i1,0],verts[i1,1],verts[i1,2])   
        p2 = vec3f(verts[i2,0],verts[i2,1],verts[i2,2]) 
        p3 = vec3f(verts[i3,0],verts[i3,1],verts[i3,2])
        points[i1] = p1
        points[i2] = p2
        points[i3] = p3  #why using face to assign verts' coordinate ??
        ebase = 3 * i
        edges[ebase].a = i1
        edges[ebase].b = i2
        edges[ebase + 1].a = i1
        edges[ebase + 1].b = i3
        edges[ebase + 2].a = i3
        edges[ebase + 2].b = i2 #compute edges

        # smoothly blend face normals into vertex normals
        p2.vec_sub(p1)
        normal = cross(p2.vec_sub(p1),p3.vec_sub(p1))
        normals[i1] = lerp(normals[i1],normal, 1.0 / (counts[i1] + 1.0))
        normals[i2] = lerp(normals[i2],normal, 1.0 / (counts[i2] + 1.0))
        normals[i3] = lerp(normals[i3],normal, 1.0 / (counts[i3] + 1.0))
        counts[i1]+= 1
        counts[i2]+= 1
        counts[i3]+= 1  

    for i in range(edgeCount):
        a = edges[i].a
        b = edges[i].b

        n1 = normals[a]
        n2 = normals[b]
        p1 = points[a]
        p2 = points[b]

        dx = p2.x - p1.x
        dy = p2.y - p1.y
        dz = p2.z - p1.z
        dd = ( dx * dx + dy * dy + dz * dz )**0.5

        dx /= dd
        dy /= dd
        dz /= dd
        dot = n1.x * n2.x + n1.y * n2.y + n1.z * n2.z
        dot2 = n2.x * dx + n2.y * dy + n2.z * dz
        ww = 1.0 - dot
        if dot2 > 0: # make it much less of a problem if convex regions have normal difference
            ww = ww ** 2
        edges[i].w = ww

        
    # Segment !
    u = segment_graph(vertexCount, edgeCount, edges, kthr)

    # Joining small segments
    for j in range(edgeCount):
        a = u.find(edges[j].a)
        b = u.find(edges[j].b)
        if (a!=b) and ((u.get_size(a) < segMinVerts) or (u.get_size(b) < segMinVerts)):
            u.join(a,b)

    # Return segment indices as vector
    outComps = np.zeros((vertexCount,),dtype=int)
    for q in range(vertexCount):
        outComps[q] = u.find(q)

    return outComps    
            

def segment_mesh(vertices,faces,kthr,segMinVerts):
    vertices = vertices.numpy()
    faces = faces.numpy()
    faces = faces.astype(int)
    vertexCount = vertices.shape[0]
    faceCount = faces.shape[0]

    comps = segment_mesh_kernel(vertices,
                                vertexCount,  
                                faces,
                                faceCount,
                                kthr,
                                segMinVerts)
    return comps                            
    
# if __name__ == "__main__":
#     faces = torch.load('/home/data1/smn/SPFormer/data/scannetv2/data_temp/face_tensor.pt')
#     vertices = torch.load('/home/data1/smn/SPFormer/data/scannetv2/data_temp/vertices_tensor.pt')

    
#     faces_len = 19429
#     vertices_len = 11103
#     faces = faces[:faces_len,:]
#     vertices = vertices[:vertices_len,:]

#     superpoint = segment_mesh(vertices, faces, kthr=0.01, segMinVerts=20)
#     pass






