from random import random
from PIL import Image
from math_function import rgb,gray,threshold_f


       
#create 8 edge
def create_edge_rgb_neghiboor(graph,img, width, x, y, x1, y1,height):
    
    vertex_id = lambda x, y: y * width + x
    

    if x > 0:
        r,g,b,sqrtRGB= rgb(img, x, y, x1-1, y1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1),r,g,b,sqrtRGB])
    if y > 0:
        r,g,b,sqrtRGB= rgb(img, x, y, x1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1, y1-1),r,g,b,sqrtRGB])
        
    if x <width-1 and y >0: 
        r,g,b,sqrtRGB= rgb(img, x, y,x1+1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1+1, y1-1),r,g,b,sqrtRGB])
       

    if x > 0 and y > 0:
        r,g,b,sqrtRGB= rgb(img, x, y, x1-1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1-1),r,g,b,sqrtRGB])
        
    if x > 0 and y <height-1:
        r,g,b,sqrtRGB= rgb(img, x, y, x1-1, y1+1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1+1),r,g,b,sqrtRGB])
        
 
    if x < width-1: 
        r,g,b,sqrtRGB= rgb(img, x, y, x1+1, y1)
        graph.append([vertex_id(x, y), vertex_id(x1+1, y1),r,g,b,sqrtRGB])
        
 
    if x <width-1 and y <height-1:
        r,g,b,sqrtRGB= rgb(img, x, y, x1+1, y1+1)
        graph.append([vertex_id(x, y), vertex_id( x1+1, y1+1),r,g,b,sqrtRGB])
        
    if y < height-1:
        r,g,b,sqrtRGB= rgb(img, x, y, x1, y1+1)
        graph.append([vertex_id(x, y), vertex_id(x1, y1+1),r,g,b,sqrtRGB])
        
 

    return graph



#create 8 edge

def create_edge_gray_neghiboor(graph,img, width, x, y, x1, y1,height):
    
    vertex_id = lambda x, y: y * width + x
   
    if x > 0:
        w = gray(img, x, y, x1-1, y1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1), w])
    if y > 0:
        w = gray(img, x, y, x1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1, y1-1), w])

    if x <width-1 and y >0:
        w = gray(img, x, y, x1+1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1+1, y1-1), w])


    if x > 0 and y > 0:
        w = gray(img, x, y, x1-1, y1-1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1-1), w])
        
    if x > 0 and y <height-1:
        w= gray(img, x, y, x1-1, y1+1)
        graph.append([vertex_id(x, y), vertex_id(x1-1, y1+1),w])
        
 
    if x < width-1: 
        w= gray(img, x, y, x1+1, y1)
        graph.append([vertex_id(x, y), vertex_id(x1+1, y1),w])
        
 
    if x <width-1 and y <height-1:
        w= gray(img, x, y, x1+1, y1+1)
        graph.append([vertex_id(x, y), vertex_id( x1+1, y1+1),w])
        
    if y < height-1:
        w= gray(img, x, y, x1, y1+1)
        graph.append([vertex_id(x, y), vertex_id(x1, y1+1),w])
        

    return graph



def build_graph(img,size,mode):
    height=size[0]
    width=size[1]
    graph = []   #new graph(empty)
    
    
    if mode == 'RGB':   # if RGB each edge have 4 weights else just 1 weghit
        for y in range(height):
            for x in range(width):
               graph= create_edge_rgb_neghiboor(graph,img, width, x, y, x-1, y,height)
    
    else:  
    
        for y in range(height):
            for x in range(width):
               graph=create_edge_gray_neghiboor(graph,img, width, x, y, x-1, y,height)
     

    return graph
    



def merege_small_components(forest, graph, min_size,size):
   
    print (' components Number: %d' % forest.conponents_size)

    image_simi_final=convolve_to_image_default(forest, size)
    image_simi_final.show()
    image_simi_final.save("output_image/imge.png")
 
    print ("image after segmentation and before removing small components (min_size)  \n")
    
    for edge in graph:
        
        
        a = forest.search_parent(edge[0])
        b = forest.search_parent(edge[1])

        if  (forest.size(a) < min_size or forest.size(b) < min_size) and a != b :
            forest.merge_nodes(a, b)
            
        
        
    print ('components Number after removing small components: %d' % forest.conponents_size)

    return convolve_to_image_default(forest, size)


    
def convolve_to_image_UserPaint(forest,size):
    
    width=size[1]
    height=size[0]    
    
    random_color = lambda: (int(random()*250), int(random()*250), int(random()*250))
   
    colors = [random_color() for i in range(width*height)]

    img = Image.new('RGB', (width, height))
    im = img.load()
    for y in range(height):
        for x in range(width):
            comp = forest.search_parent(y * width + x)
            comp=forest.nodes[comp].parent_component
            im[x, y] = colors[comp]

    return img.transpose(Image.ROTATE_270).transpose(Image.FLIP_LEFT_RIGHT)
    
 

def convolve_to_image_default(forest,size):
    
    width=size[1]
    height=size[0]    
    
    random_color = lambda: (int(random()*254), int(random()*254), int(random()*254))
   
    colors = [random_color() for i in range(width*height)]

    img = Image.new('RGB', (width, height))
    im = img.load()
    for y in range(height):
        for x in range(width):
            comp = forest.search_parent(y * width + x)
            im[x, y] = colors[comp]

    return img.transpose(Image.ROTATE_270).transpose(Image.FLIP_LEFT_RIGHT)
    
 
  
    #the main function to segmnet an givin image
def segment_image(smooth_img,size, const,mode,min_size,algorithm1):
    
    # first bulid a graph that contain edge for every neighbor node
    graph = build_graph(smooth_img, size,mode)
    
    if(mode == 'RGB'):
        if(algorithm1):
            weight_r = lambda edge: edge[2]  # wight for R color
            weight_g = lambda edge: edge[3]   # wight for G color
            weight_b = lambda edge: edge[4]   # wight for B color
           
            weight=lambda edge: edge[5]
            nodes_size=size[1]*size[0]
           
            forest1 = Forest(nodes_size)
            forest2= Forest(nodes_size)
            forest3 = Forest(nodes_size)
            forest= Forest(nodes_size)
    
            sorted_graph1 = sorted(graph, key=weight_r)
            sorted_graph2 = sorted(graph, key=weight_g)
            sorted_graph3 = sorted(graph, key=weight_b)
            sorted_graph = sorted(graph, key=weight)
    
    
            threshold1 = [threshold_f(1, const)] * nodes_size
            threshold2 = [threshold_f(1, const)] * nodes_size
            threshold3 = [threshold_f(1, const)] * nodes_size
    
            for edge in sorted_graph1:
                left_node = forest1.search_parent(edge[0])
                right_node = forest1.search_parent(edge[1])
           
                if left_node != right_node and weight_r(edge) <= threshold1[left_node] and weight_r(edge) <= threshold1[right_node] :
                    forest1.merge_nodes(left_node, right_node)
                    a = forest1.search_parent(left_node)
                    threshold1[a] = weight_r(edge) + threshold_f(forest1.nodes[a].size, const)
           
            image_R= merege_small_components(forest1, sorted_graph1, min_size,size)
            image_R.show()      
            image_R.save("output_image/forest_R.png")

            for edge in sorted_graph2:
                left_node = forest2.search_parent(edge[0])
                right_node = forest2.search_parent(edge[1])
      
      
                if left_node != right_node and weight_b(edge) <= threshold2[left_node] and weight_b(edge) <= threshold2[right_node] :
                    forest2.merge_nodes(left_node, right_node)
                    a = forest2.search_parent(left_node)
                    threshold2[a] = weight_b(edge) + threshold_f(forest2.nodes[a].size, const)
           
            image_B= merege_small_components(forest2, sorted_graph2, min_size,size)
            image_B.show()      
            image_B.save("output_image/forest_B.png")

            for edge in sorted_graph3:
                left_node = forest3.search_parent(edge[0])
                right_node = forest3.search_parent(edge[1])

                if left_node != right_node and  weight_g(edge) <= threshold3[left_node] and weight_g(edge) <= threshold3[right_node] :
                    forest3.merge_nodes(left_node, right_node)
                    a = forest3.search_parent(left_node)
                    threshold3[a] = weight_g(edge) + threshold_f(forest3.nodes[a].size, const)
                    
            image_G= merege_small_components(forest3, sorted_graph3, min_size,size)
            image_G.show() 
            image_G.save("output_image/forest_G.png")
            
            
            
            
            # merge 3 forest R,G,B
            
            
            for edge in sorted_graph:
                
                
                left_node = forest.search_parent(edge[0])
                right_node = forest.search_parent(edge[1])
                
                left_node_r = forest1.search_parent(edge[0])
                right_node_r = forest1.search_parent(edge[1])
                
                left_node_b = forest2.search_parent(edge[0])
                right_node_b = forest2.search_parent(edge[1])
                
                left_node_g = forest3.search_parent(edge[0])
                right_node_g = forest3.search_parent(edge[1])
              
                rgb_condition = (left_node_r ==right_node_r) and (left_node_b ==right_node_b) and (left_node_g ==right_node_g)
               
              
                if left_node != right_node and rgb_condition  :
                    forest.merge_nodes(left_node, right_node)
        
        
        
        
        else :
             weight = lambda edge: edge[5]
             nodes_size=size[1]*size[0]
             forest = Forest(nodes_size)
             sorted_graph = sorted(graph, key=weight)
             threshold = [threshold_f(1, const)] * nodes_size
             for edge in sorted_graph:
                left_node = forest.search_parent(edge[0])
                right_node = forest.search_parent(edge[1])


                if left_node != right_node and weight(edge) <= threshold[left_node] and weight(edge) <= threshold[right_node] :
                    forest.merge_nodes(left_node, right_node)
                    a = forest.search_parent(left_node)
                    threshold[a] = weight(edge) + threshold_f(forest.nodes[a].size, const)
           
  
    
    else:
    #Gray Image
        
        weight = lambda edge: edge[2]
        nodes_size=size[1]*size[0]
        forest = Forest(nodes_size)
        sorted_graph = sorted(graph, key=weight)
        threshold = [threshold_f(1, const)] * nodes_size
    
        for edge in sorted_graph:
            left_node = forest.search_parent(edge[0])
            right_node = forest.search_parent(edge[1])


            if (left_node != right_node and weight(edge) <= threshold[left_node] and weight(edge) <= threshold[right_node]):
                forest.merge_nodes(left_node, right_node)
                a = forest.search_parent(left_node)
                threshold[a] = weight(edge) + threshold_f(forest.nodes[a].size, const)
            
        
   
    image_final=convolve_to_image_default(forest, size)
    result=image_final,forest,sorted_graph
    return result


#data structure

class Pixel:
    def __init__(self, parent):
        
        self.parent_component=parent
        self.rank = 0
        self.parent = parent
        self.size = 1

 
class Forest:
    
    def __init__(self, nodes_size):
        self.nodes = [Pixel(i) for i in range(nodes_size)]
        self.conponents_size = nodes_size
  
    def search_parent(self, n):
        temp = n
        while temp != self.nodes[temp].parent:
            temp = self.nodes[temp].parent

        self.nodes[n].parent = temp
        return temp


    def size(self, i):
        return self.nodes[i].size


    def merge_nodes(self, a, b):
        if self.nodes[a].rank > self.nodes[b].rank:
            
            self.nodes[a].parent_component=b
            self.nodes[b].parent = a
            self.nodes[a].size = self.nodes[a].size + self.nodes[b].size
        else:
            
            self.nodes[a].parent = b
            self.nodes[b].size = self.nodes[b].size + self.nodes[a].size
            self.nodes[a].parent_component=b
            if self.nodes[a].rank == self.nodes[b].rank:
                self.nodes[b].rank = self.nodes[b].rank + 1

        self.conponents_size = self.conponents_size - 1
 



