# -*- coding: utf-8 -*-
"""
Created on Tuesday July 16 11:30:00 2019
E-mail = yuyafang1014@163.com
@author: Yuyafang
"""

import datetime
import sys

import cv2
import matplotlib.pyplot as plt
import numpy as np
from osgeo import gdal
from skimage import segmentation
from skimage.future import graph
from skimage.segmentation import mark_boundaries
from skimage.segmentation import slic


def _weight_mean_color(graph, src, dst, n):
    """
    计算两个节点的均值颜色之间的差异，用于构建图的权重。
    :param graph:
    :param src:
    :param dst:
    :param n:
    :return:
    """
    diff = graph.node[dst]['mean color'] - graph.node[n]['mean color']
    diff = np.linalg.norm(diff)
    return {'weight': diff}


def merge_mean_color(graph, src, dst):
    """
    将两个节点的颜色信息合并为一个节点的均值颜色。
    :param graph:
    :param src:
    :param dst:
    :return:
    """
    graph.node[dst]['total color'] += graph.node[src]['total color']
    graph.node[dst]['pixel count'] += graph.node[src]['pixel count']
    graph.node[dst]['mean color'] = (graph.node[dst]['total color'] /
                                     graph.node[dst]['pixel count'])


def daw_figure(num_segment, image, segment):
    """
    绘制 SLIC 分割的结果。将SLIC分割结果可视化为带有超像素边界的图像。
    :param num_segment:超像素个数
    :param image:从.tif图像中取出的r、g、b三通道的图像
    :param segment:超像素分割后的结果
    """
    fig = plt.figure("slic-superpixels -- %d segments" % num_segment)
    ax = fig.add_subplot(1, 1, 1)
    ax.imshow(mark_boundaries(image, segment))
    plt.axis("off")
    return 0


def main():
    # ------Slic-Superpixel segmentation------

    # 命令行示例：python slic-superpixel.py input_image.tif 3500 5 20
    # 检查程序运行时命令行参数的数量，以确保用户提供了足够的参数来运行程序
    if len(sys.argv) < 4:
        print("need more args")
        # TODO: print usage
        exit(-1)

    # Load the image(GDAL)
    dataset = gdal.Open(sys.argv[1])  # input_image.tif
    super_pixel_size = int(sys.argv[2])  # 3500
    """是超像素的尺寸，即一个超像素占多少个像素。不是超像素的个数"""

    nc_sigma = int(sys.argv[3])  # 5
    """图像每个维度进行预处理时的高斯平滑核宽。若给定为标量值，则同一个值运用到各个维度。0意味
    着不平滑。如果“sigma”是标量的，并且提供了手动体素间距，则自动缩放它（参见注释部分）"""

    scale = int(sys.argv[4])  # 20

    im_width = dataset.RasterXSize
    im_height = dataset.RasterYSize
    im_bands = dataset.RasterCount

    if im_bands == 3:
        band1 = dataset.GetRasterBand(1)
        band2 = dataset.GetRasterBand(2)
        band3 = dataset.GetRasterBand(3)
        im_red_band = band1.ReadAsArray(0, 0, im_width, im_height)
        im_green_band = band2.ReadAsArray(0, 0, im_width, im_height)
        im_blue_band = band3.ReadAsArray(0, 0, im_width, im_height)
        merge_img = cv2.merge([im_red_band, im_green_band, im_blue_band])
        """从.tif图像中取出的r、g、b三通道的图像"""

    # 得到超像素的个数。可能最后分割出的块数与实际设置并不一样，可能是slic算法做了后续处理，将小的超像素合并到大的超像素中。
    num_segments = int(im_width * im_height / super_pixel_size)
    """超像素个数"""

    start1 = datetime.datetime.now()

    # Apply SLIC and extract (approximately) the supplied number of segments
    segments = slic(merge_img, n_segments=num_segments, sigma=nc_sigma)
    """
    超像素分割后的结果，它是一个与原始图像大小相同的二维数组，每个像素都被分配一个超像素标签。
    是一个存储了超像素分割后的标签的数据结构。
    它包含了超像素分割后每个像素的标签。
    segments[y, x]将返回给定坐标 (x, y) 处像素所属的超像素标签
    """

    print(segments)
    print(segments.shape)

    cv2.waitKey(0)

    # Show the output of SLIC
    end1 = datetime.datetime.now()
    print("The time spent in segmentation (seconds): ", (end1 - start1).seconds)
    daw_figure(num_segments, merge_img, segments)

    # Show the plots
    plt.show()

    # 在SLIC超像素分割之后，程序构建了一个区域邻接图 (Region Adjacency Graph, RAG)，然后使用分层的方式对这些区域进行合并。

    # ------Hierarchical Region Merging Algorithm for Region Adjacent Graph (RAG)------
    start2 = datetime.datetime.now()

    # Compute the Region Adjacency Graph using mean colors
    g = graph.rag_mean_color(merge_img, segments)
    """
    构建的一个区域邻接图 (RAG)，使用均值颜色信息来表示超像素之间的关系。
    """

    end21 = datetime.datetime.now()
    print("The time spent in mean colors (seconds): ", (end21 - start2).seconds)

    # Perform hierarchical merging of a RAG
    labels = graph.merge_hierarchical(segments, g, thresh=scale, rag_copy=False,
                                      in_place_merge=True,
                                      merge_func=merge_mean_color,
                                      weight_func=_weight_mean_color)
    """
    包含了区域合并后的标签，这些标签标识了被合并的区域。
    使用分层的方式对这些区域进行合并，通过调用 graph.merge_hierarchical 函数，指定合并阈值和合并函数。
    labels 并不包含每个像素的超像素标签，而是包含了超像素之间的合并标签。这个变量用于存储分层区域合并后的结果，而不是像素级别的标签
    labels 存储了分层区域合并后的结果，其中每个标签代表一个合并后的超像素区域。这些标签对应于区域的标识符，而不是像素级别的标签。
    """

    # Return image with boundaries between labeled regions highlighted
    out = segmentation.mark_boundaries(segments, labels)
    """
    在 out 图像中，超像素之间的边界被标记为边缘线，以突出显示超像素之间的边界。
    这可以更清晰地看到超像素的分割结果，以及合并后的区域之间的边界。
    """

    # Show the figure
    end22 = datetime.datetime.now()
    fig = plt.figure("Region Merging -- %d scale" % (scale))
    ax = fig.add_subplot(1, 1, 1)
    print("The time spent on region merging (seconds): ", (end22 - start2).seconds)
    ax.imshow(out)
    plt.axis("off")
    plt.show()

    # ------Draw equalized histogram------
    start3 = datetime.datetime.now()
    arr = segments.flatten()

    # Show equalized histogram
    plt.figure("Equalized histogram")
    plt.hist(arr, bins=256, normed=1, edgecolor='None', facecolor='red')
    end3 = datetime.datetime.now()
    print("The time spent on equalized histogram (seconds): ", (end3 - start3).seconds)
    plt.show()


'''
usage:  python ./slic-superpixel.py input_image.tif super_pixel_size sigma scale
example: python ./slic-superpixel.py a.tif 3500 5 20

'''
if __name__ == '__main__':
    main()