def test_minsize():
    # single-channel:
    img = data.coins()[20:168, 0:128]
    for min_size in np.arange(10, 100, 10):
        segments = felzenszwalb(img, min_size=min_size, sigma=3)
        counts = np.bincount(segments.ravel())
        # actually want to test greater or equal.
        assert_greater(counts.min() + 1, min_size)
    # multi-channel:
    coffee = data.coffee()[::4, ::4]
    for min_size in np.arange(10, 100, 10):
        segments = felzenszwalb(coffee, min_size=min_size, sigma=3)
        counts = np.bincount(segments.ravel())
        # actually want to test greater or equal.
        assert_greater(counts.min() + 1, min_size)

def fun_compare_colorsegmentation_and_display(image_data, number_segments=250, compactness_factor=10):
    """
    The function is a copy of what does this link http://scikit-image.org/docs/dev/auto_examples/plot_segmentations.html
    """
    segments_fz = felzenszwalb(image_data, scale=100, sigma=0.5, min_size=50)
    segments_slic = slic(image_data, n_segments=number_segments, compactness=compactness_factor, sigma=1)
    segments_quick = quickshift(image_data, kernel_size=3, max_dist=6, ratio=0.5)

    print ("Felzenszwalb's number of segments: %d" % len(np.unique(segments_fz)))
    print ("Slic number of segments: %d" % len(np.unique(segments_slic)))
    print ("Quickshift number of segments: %d" % len(np.unique(segments_quick)))

    fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, subplot_kw={"adjustable": "box-forced"})
    fig.set_size_inches(8, 3, forward=True)
    fig.subplots_adjust(0.05, 0.05, 0.95, 0.95, 0.05, 0.05)

    ax[0].imshow(mark_boundaries(image_data, segments_fz, color=(1, 0, 0)))
    ax[0].set_title("Felzenszwalbs's method")
    ax[1].imshow(mark_boundaries(image_data, segments_slic, color=(1, 0, 0)))
    ax[1].set_title("SLIC")
    ax[2].imshow(mark_boundaries(image_data, segments_quick, color=(1, 0, 0)))
    ax[2].set_title("Quickshift")
    for a in ax:
        a.set_xticks(())
        a.set_yticks(())
    plt.show()

def selectiveSearch(image):
    segments = felzenszwalb(image, scale=kFelzenszwalbScale)
    numRegions = segments.max()

    rectangles = []

    for regionTag in range(numRegions):
        selectedRegion = segments == regionTag
        regionPixelIndices = np.transpose(np.nonzero(selectedRegion))
        rectangle = aabb(regionPixelIndices)
        rectangles.append(rectangle)

    # Implement similarities, neighbourhood merging.
    # Felzenszwalb's segmentation is ridiculously good already.

    def debug():
        marked = np.zeros(image.shape, dtype=np.uint8)

        for rectangle in rectangles:
            rr, cc = rectangle.pixels(marked.shape)
            randcolor = randint(0, 255), randint(0, 255), randint(0, 255)
            marked[rr, cc] = randcolor

        print(image.shape, segments.shape, marked.shape)

        io.imshow_collection([image, segments, marked])
        io.show()

    # debug()

    return rectangles

def SegmentationFelz_run_2d(rod):
    img = img_as_float(
        RodriguesToUnambiguousColor(rod["x"], rod["y"], rod["z"], maxRange=None, centerR=None).astype("uint8")
    )
    segments_slic = felzenszwalb(img, scale=100, sigma=0.0, min_size=10)
    print("Slic number of segments: %d" % len(np.unique(segments_slic)))
    return segments_slic

def mask_felz(image, config):
    #constants for felzenszwalb segmentation function
    scale = config[':felzenszwalb'][':scale']
    sigma = config[':felzenszwalb'][':sigma'] 
    min_size = config[':felzenszwalb'][':min_size'] 

    segments = felzenszwalb(image, scale, sigma, min_size)
    return segments

def doSegment(param, img):
  if param[0] == 'slic':
    segments_res = slic(img, n_segments=int(param[1]), compactness=int(param[2]), sigma=int(param[3]), convert2lab=True)
  elif param[0] == 'pff':
    segments_res = felzenszwalb(img, scale=int(param[1]), sigma=float(param[2]), min_size=int(param[3]))
  elif param[0] == 'quick':
    segments_res = quickshift(img, kernel_size=int(param[1]), max_dist=int(param[2]), ratio=float(param[3]), convert2lab=True)
  return segments_res

def test_minsize():
    # single-channel:
    img = data.coins()[20:168,0:128]
    for min_size in np.arange(10, 100, 10):
        segments = felzenszwalb(img, min_size=min_size, sigma=3)
        counts = np.bincount(segments.ravel())
        # actually want to test greater or equal.
        assert_greater(counts.min() + 1, min_size)
    # multi-channel:
    coffee = data.coffee()[::4, ::4]
    for min_size in np.arange(10, 100, 10):
        segments = felzenszwalb(coffee, min_size=min_size, sigma=3)
        counts = np.bincount(segments.ravel())
        # actually want to test greater or equal.
        # the construction doesn't guarantee min_size is respected
        # after intersecting the sementations for the colors
        assert_greater(np.mean(counts) + 1, min_size)

def test_merging():
    # test region merging in the post-processing step
    img = np.array([[0, 0.3], [0.7, 1]])
    # With scale=0, only the post-processing is performed.
    seg = felzenszwalb(img, scale=0, sigma=0, min_size=2)
    # we expect 2 segments:
    assert_equal(len(np.unique(seg)), 2)
    assert_array_equal(seg[0, :], 0)
    assert_array_equal(seg[1, :], 1)

    def __init__(self, fname):
        self.fname = fname
        self.image_cache = {}

        print 'loading image'
        self.raw_img = io.imread(fname)
        self._segment_ids = None
        img_float = img_as_float(self.raw_img)
        print 'segmenting image'
        self.segments = felzenszwalb(img_float, scale=300, sigma=0.5, min_size=100)

 def extract(self, image):
     """
     Performs segmentation and returns a set of nd.array
     """
     # Init the list of segmentations
     segmentations = []
     for param in self.params_:
         sc, sg, m = param
         segm_mask = segmentation.felzenszwalb(image, sc, sg, m)
         segmentations.append(segm_mask)
     return segmentations

def test_grey():
    # very weak tests.
    img = np.zeros((20, 21))
    img[:10, 10:] = 0.2
    img[10:, :10] = 0.4
    img[10:, 10:] = 0.6
    seg = felzenszwalb(img, sigma=0)
    # we expect 4 segments:
    assert_equal(len(np.unique(seg)), 4)
    # that mostly respect the 4 regions:
    for i in range(4):
        hist = np.histogram(img[seg == i], bins=[0, 0.1, 0.3, 0.5, 1])[0]
        assert_greater(hist[i], 40)

def test_color():
    # very weak tests.
    img = np.zeros((20, 21, 3))
    img[:10, :10, 0] = 1
    img[10:, :10, 1] = 1
    img[10:, 10:, 2] = 1
    seg = felzenszwalb(img, sigma=0)
    # we expect 4 segments:
    assert_equal(len(np.unique(seg)), 4)
    assert_array_equal(seg[:10, :10], 0)
    assert_array_equal(seg[10:, :10], 2)
    assert_array_equal(seg[:10, 10:], 1)
    assert_array_equal(seg[10:, 10:], 3)

def SegmentationFelzenszwalb(filename="/media/scratch/plasticity/lvp2d1024_s0_d.tar", time=None):
    t, s = TarFile.LoadTarState(filename, time=time)
    rod = s.CalculateRotationRodrigues()
    img = img_as_float(
        RodriguesToUnambiguousColor(rod["x"], rod["y"], rod["z"], maxRange=None, centerR=None).astype("uint8")
    )

    segments_slic = felzenszwalb(img, scale=100, sigma=0.0, min_size=10)
    print("Slic number of segments: %d" % len(np.unique(segments_slic)))

    fig = plt.figure()
    plt.imshow(mark_boundaries(img, segments_slic, color=[0, 0, 0], outline_color=None))
    plt.show()

    def __init__(self, fname):
        self.fname = fname

        print 'loading image'
        img = Image.open(fname)
        img.thumbnail((1280,1024), Image.ANTIALIAS)
        width, height = img.size
        square_size = width * height
        min_size = square_size / 300
        self.raw_img = np.array(enhance(img).convert('RGB'))
        self._segment_ids = None
        img_float = img_as_float(self.raw_img)
        print 'segmenting image'
        self.segments = felzenszwalb(img_float, scale=300, sigma=0.25, min_size=min_size)

def clustering(inPixNP, width, height, scale=50, sigma=4.5, min_size=10):
    segmentsNP = felzenszwalb(inPixNP, scale, sigma, min_size)
    # felzenszwalb(image, scale=1, sigma=0.8, min_size=20)
    # image : (width, height, 3) or (width, height) ndarray - Input image.
    # scale : float - Free parameter. Higher means larger clusters.
    # sigma : float - Width of Gaussian kernel used in preprocessing.
    # min_size : int - Minimum component size. Enforced using postprocessing.
    
    # create a data structure with regionIDs as keys and lists of their pixel coordinates as values:
    maxSegmentID = np.max(segmentsNP)
    segmentsByIDlist = [[] for i in range(maxSegmentID + 1)]
    for y in range(0,height):
        for x in range(0,width):
            regionID = segmentsNP[y,x]
            segmentsByIDlist[regionID].append((x, y))
    return segmentsNP, segmentsByIDlist

def demo(image_name,color_space_list=None,ks=None,sim_feats_list=None,net='vgg16', cpu_mode=True):
	''' Object Recognition Demo : Selective Search + RCNN
	parameters
	----------
	image_name : filename of image stored in 'Data/img'
	
	color_space_list : list of colorspaces to be used. Refer color_utils for list of possible colorspaces.
	Default : [ 'HSV', 'LAB']

	ks : list felzenszwalb scale/threshold and minimum segment size.
	Default : [50, 100]
	
	'''
	
	blob_array = []	
	priority = []
	img = plt.imread('Data/img/' + image_name + '.jpg')
	seg_dir = 'Data/segments/'
	if color_space_list is None: color_space_list = ['HSV','LAB']
	if ks is None: ks = [50,100]
	if sim_feats_list is None: sim_feats_list = [[ sf.color_hist_sim(), sf.texture_hist_sim(), sf.size_sim(img.shape), sf.fill_sim(img.shape) ],[ sf.texture_hist_sim(), sf.size_sim(img.shape), sf.fill_sim(img.shape) ]]

	cc = convert_colorspace(img,color_space_list)
	seg_filename = [seg_dir + 'HSV/50/' + image_name +'.mat',seg_dir + 'HSV/100/' + image_name +'.mat', seg_dir + 'LAB/50/' + image_name +'.mat',seg_dir + 'LAB/100/' + image_name +'.mat']

	for i in range(len(color_space_list)):
		for j in range(len(ks)):
			for k in range(len(sim_feats_list)):
				_img = cc[i]
				_file = "%s%s/%d/%s.mat"%(seg_dir,color_space_list[i].upper(),ks[j],image_name)
				if not os.path.exists(_file):
					segment_mask = felzenszwalb(_img,scale=ks[j],sigma=0.8,min_size=ks[j])
					_temp_dict = dict()
					_temp_dict['blobIndIm'] = segment_mask + 1
					scipy.io.savemat(_file,_temp_dict)
				_blob_array = ssearch._ssearch(_img,ssearch.load_segment_mask(_file),sim_feats = sim_feats_list[k])
				blob_array.append(_blob_array)
				priority.append( np.arange(len(_blob_array),0,-1).clip(0,(len(_blob_array)+1)/2))
		
	bboxes = ssearch.remove_duplicate(blob_array,priority)
	bbox_dict = {}
	bbox_dict['boxes'] = np.vstack([np.asarray(bboxes)[:,2],np.asarray(bboxes)[:,1],np.asarray(bboxes)[:,4],np.asarray(bboxes)[:,3]]).T
	print('\nComputed %d proposals'%(len(bboxes)))
	scipy.io.savemat('Data/Boxes/' + image_name + '.mat',bbox_dict)
	rcnn.rcnn_demo(image_name,net=net, cpu_mode=cpu_mode)

def FELZENSZWALB(Input_Image, scale, sigma, min_size):
   
    '''
    Description:   Computes Felsenszwalbs efficient graph based image segmentation. 

    source:         skimage, openCv python 
    
    parameters:     Input_Image : ndarray
                    Input image
    
                    min-size : int
                    Minimum component size. Enforced using postprocessing.
    
                    scale:  float
                     The parameter scale sets an observation level. Higher scale means less and larger segments.
    
                    sigma : float
                    Width of Gaussian smoothing kernel for preprocessing. Zero means no smoothing.
    
    return:         Segment_mask : ndarray
                    Integer mask indicating segment labels.
        
    '''
    #default values, set in case of 0 as input
    if scale == 0:
        scale = 5
    if sigma == 0:
        sigma = 0.5
    if min_size == 0:
        min_size = 30

    #print Input
    img = cv2.imread(Input_Image)
    #print img
    #print img.shape

    segments_fz = felzenszwalb(img, scale, sigma, min_size)
    print segments_fz.shape
    #print ('segments_fz datatype',segments_fz.dtype )
    print("Felzenszwalb's number of segments: %d" % len(np.unique(segments_fz)))
    print ('segments_fz datatype',segments_fz.dtype )

    return segments_fz

def doSegment(im):
    timestr = time.strftime("%Y%m%d-%H%M%S")
    global count

    img = img_as_float(im)
    # img = img_as_float(imgO[::2, ::2])

    segments_fz = felzenszwalb(img, scale=2.0, sigma=1.0, min_size=0)
    # segments_slic = slic(img, n_segments=250, compactness=10, sigma=1)
    # segments_quick = quickshift(img, kernel_size=5, max_dist=10, ratio=0.0)

    if count % 2 == 0:
        bounds = mark_boundaries(img, segments_fz, color=(0.0, 0.0, 0.0), mode="inner")
    else:
        bounds = mark_boundaries(img, segments_fz, color=(1, 1, 1), mode="inner")

    labeled = label2rgb(
        segments_fz,
        img,
        alpha=1.0,
        bg_color=(0, 0, 0),
        bg_label=0,
        colors=[
            (0, 0, 0),
            (1, 1, 1),
            (0.25, 0.25, 0.25),
            (0.5, 0.5, 0.5),
            (0.75, 0.75, 0.75),
            (0.1, 0.1, 0.1),
            (0.05, 0.05, 0.05),
            (0.2, 0.2, 0.2),
            (0.15, 0.15, 0.15),
        ],
    )

    bb = Image.fromarray(np.uint8(labeled * 255))
    bb.show()

    bb.save("output/white" + timestr + ".png")
    # bb = bb.resize((500, 500), Image.NEAREST)
    count = count + 1

def felzenszwalb(request):

    image, response = process(request, color = True)

    if len(image.shape) == 2:
        
        image = cv2.cvtColor(image,cv2.COLOR_GRAY2RGB)

    scale = int(request.GET.get('felzenszwalb-scale'))

    sigma = float(request.GET.get('felzenszwalb-sigma'))

    min_size = int(request.GET.get('felzenszwalb-min'))  

    segments = segmentation.felzenszwalb(image, scale = scale, sigma = sigma, min_size = min_size)

    image = segmentation.mark_boundaries(image, segments)

    io.imsave(response['filename'], image)

    return JsonResponse(response)

def segmentation(path):
    ds = gdal.Open(path, GA_ReadOnly )
    img = np.array(ds.GetRasterBand(1).ReadAsArray())
    segments_fz = felzenszwalb(img, scale=100, sigma=0.5, min_size=50)
    #segments_slic = slic(img, n_segments=250, compactness=10, sigma=1)
    #segments_quick = quickshift(img, kernel_size=3, max_dist=6, ratio=0.5)

    print("Felzenszwalb's number of segments: %d" % len(np.unique(segments_fz)))

    fig, ax = plt.subplots(1, 3)
    fig.set_size_inches(8, 3, forward=True)
    fig.subplots_adjust(0.05, 0.05, 0.95, 0.95, 0.05, 0.05)

    ax[0].imshow(mark_boundaries(img, segments_fz))
    ax[0].set_title("Felzenszwalbs's method")
    # ax[1].imshow(mark_boundaries(img, segments_slic))
    # ax[1].set_title("SLIC")
    # ax[2].imshow(mark_boundaries(img, segments_quick))
    # ax[2].set_title("Quickshift")
    for a in ax:
        a.set_xticks(())
        a.set_yticks(())
    plt.show()