#!/usr/bin/python
# -*- coding: utf-8 -*-
#
"""
  Nome: Slic.py
  Autor: Alessandro dos Santos Ferreira ( santosferreira.alessandro@gmail.com )

  Descricão: TODO.
"""

import cv2
import numpy as np
from skimage.segmentation import slic
from skimage.segmentation import mark_boundaries
from skimage.util import img_as_float, img_as_ubyte
from collections import OrderedDict

from util.Config import Config
from util.Utils import TimeUtils
from util.X11Colors import Colors

from Segmentation import Segmentation

class Slic(Segmentation):
    
    n_segments = None
    sigma = None
    compactness = None
    border_color = None
    border_outline = None
    
    __segments = None
    
    __original_image = None
    
    def __init__(self, n_segments = 100, sigma = 5.0, compactness = 10.0, border_color = 'Yellow', border_outline = 'No'):
        self.n_segments = Config("Segments", n_segments, int)
        self.sigma = Config("Sigma", sigma, float)
        self.compactness = Config("Compactness", compactness, float)
        self.border_color = Config("Border Color", border_color, 'color')
        self.border_outline = Config("Border Outline", border_outline, str)
        
    
    def get_config(self):
        slic_config = OrderedDict()
        
        slic_config["n_segments"] = self.n_segments
        slic_config["sigma"] = self.sigma
        slic_config["compactness"] = self.compactness
        slic_config["border_color"] = self.border_color
        slic_config["border_outline"] = self.border_outline
        
        return slic_config
        
    def set_config(self, configs):
        self.n_segments = Config.nvl_config(configs["n_segments"], self.n_segments)
        self.sigma = Config.nvl_config(configs["sigma"], self.sigma)
        self.compactness = Config.nvl_config(configs["compactness"], self.compactness)
        self.border_color = Config.nvl_config(configs["border_color"], self.border_color)
        self.border_outline = Config.nvl_config(configs["border_outline"], self.border_outline)
        
        self.border_outline.value = self.border_outline.value if self.border_outline.value == 'Yes' else 'No'

    def get_summary_config(self):
        slic_config = OrderedDict()
        
        slic_config[self.n_segments.label] = self.n_segments.value
        slic_config[self.sigma.label] = self.sigma.value
        slic_config[self.compactness.label] = self.compactness.value
        slic_config[self.border_color.label] = self.border_color.value
        slic_config[self.border_outline.label] = self.border_outline.value
        
        summary = ''
        for config in slic_config:
            summary += "%s: %s\n" % (config, str(slic_config[config]))
        
        return summary
    

    def get_segment(self, px = 0, py = 0, idx_segment = None):
        if self.__segments is None:
            return None, 0, -1, 0
        
        start_time = TimeUtils.get_time()
        
        if idx_segment is None:
            idx_segment = self.__segments[py, px]
        
        mask_segment = np.zeros(self.__original_image.shape[:2], dtype="uint8")
        mask_segment[self.__segments == idx_segment] = 255
        size_segment = mask_segment[self.__segments == idx_segment].size

        segment = self.__original_image.copy()
        segment = cv2.bitwise_and(segment, segment, mask=mask_segment)

        contours, _  = cv2.findContours(mask_segment,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)[-2:]

        m = -1
        max_contour = None
        for cnt in contours:
            if (len(cnt) > m):
                m = len(cnt)
                max_contour = cnt

        x,y,w,h = cv2.boundingRect(max_contour)
        segment = segment[y:y+h, x:x+w]
        
        end_time = TimeUtils.get_time()
        
        return segment, size_segment, idx_segment, (end_time - start_time)
        
        
    def paint_segment(self, image, color, px = 0, py = 0, idx_segment = None, border = True, clear = False):
        
        if self.__segments is None:
            return image, 0
            
        start_time = TimeUtils.get_time()
        
        if idx_segment is None:
            idx_segment = self.__segments[py, px]
        height, width, channels = self.__original_image.shape
        
        mask_segment = np.zeros(self.__original_image.shape[:2], dtype="uint8")
        mask_segment[self.__segments == idx_segment] = 255
        mask_inv = cv2.bitwise_not(mask_segment)
            
        class_color = np.zeros((height,width,3), np.uint8)
        class_color[:, :] = Colors.get_color(color)
        if clear == False:
            colored_image = cv2.addWeighted(self.__original_image, 0.7, class_color, 0.3, 0)
        else:
            colored_image = self.__original_image
        colored_image = cv2.bitwise_and(colored_image, colored_image, mask=mask_segment)
        
        new_image = cv2.bitwise_and(image, image, mask=mask_inv)
        mask_segment[:] = 255
        new_image = cv2.bitwise_or(new_image, colored_image, mask=mask_segment)

        if border == True:
            color = Colors.get_color_zero_one(self.border_color.get_cast_val())
            outline_color = color if self.border_outline.value == 'Yes' else None
            
            new_image = img_as_ubyte( mark_boundaries(img_as_float(new_image), self.__segments.astype(np.int8), color=color, outline_color=outline_color) ) 
            
        end_time = TimeUtils.get_time()
        
        return new_image, (end_time - start_time)
        

    def run(self, image):
        
        self.__original_image = image
        
        color = Colors.get_color_zero_one(self.border_color.get_cast_val())
        outline_color = color if self.border_outline.value == 'Yes' else None
        
        start_time = TimeUtils.get_time()
        self.__segments = slic(img_as_float(image), 
                            n_segments=self.n_segments.get_cast_val(), 
                            sigma=self.sigma.get_cast_val(), 
                            compactness=self.compactness.get_cast_val())
        end_time = TimeUtils.get_time()
        
        return img_as_ubyte( mark_boundaries(image, self.__segments.astype(np.int8), color=color, outline_color=outline_color) ), (end_time - start_time)