# 去除过曝光像素
# 导入必要的库

import matplotlib

matplotlib.use('TkAgg')
from plantcv import plantcv as pcv
import get_glcm
import os
import cv2
import numpy as np
from plantcv.plantcv._helpers import _iterate_analysis, _cv2_findcontours, _object_composition, _grayscale_to_rgb, \
    _scale_size
from plantcv.plantcv import outputs, within_frame
from plantcv.plantcv import params
from plantcv.plantcv._debug import _debug
from outputs2dataframe import process_plantcv_outputs


def size(img, labeled_mask, n_labels=1, label=None):
    """A function that analyzes the shape and size of objects and outputs data.

    Inputs:
    img          = RGB or grayscale image data for plotting
    labeled_mask = Labeled mask of objects (32-bit).
    n_labels     = Total number expected individual objects (default = 1).
    label        = Optional label parameter, modifies the variable name of
                   observations recorded (default = pcv.params.sample_label).

    Returns:
    analysis_image = Diagnostic image showing measurements.

    :param img: numpy.ndarray
    :param labeled_mask: numpy.ndarray
    :param n_labels: int
    :param label: str
    :return analysis_image: numpy.ndarray
    """
    # Set lable to params.sample_label if None
    if label is None:
        label = params.sample_label

    img = _iterate_analysis(img=img, labeled_mask=labeled_mask, n_labels=n_labels, label=label, function=_analyze_size)
    # Debugging
    _debug(visual=img, filename=os.path.join(params.debug_outdir, str(params.device) + '_shapes.png'))
    return img


def _analyze_size(img, mask, label):
    """Analyze the size and shape of individual objects."""

    params.device += 1
    # Initialize analysis output values
    area = 0
    hull_area = 0
    solidity = 0
    perimeter = 0
    width = 0
    height = 0
    caliper_length = 0
    longest_path = 0
    cmx = 0
    cmy = 0
    hull_vertices = 0
    circularity = 0
    shape_factor = 0
    aspect_ratio = 0
    moments = {}

    # Check if the object is touching image boundaries (QC)
    in_bounds = within_frame(mask=mask, label=label)

    # Convert grayscale images to color
    img = _grayscale_to_rgb(img)
    # Plot image
    plt_img = np.copy(img)

    # Find contours
    cnt, cnt_str = _cv2_findcontours(bin_img=mask)

    # Consolidate contours
    obj = _object_composition(contours=cnt, hierarchy=cnt_str)
    contour_list = obj.squeeze().tolist()
    # Analyze shape properties if the object is large enough
    if len(obj) > 5:
        # Convex Hull
        hull = cv2.convexHull(obj)
        # Number of convex hull vertices
        hull_vertices = len(hull)
        # Convex Hull area
        hull_area = cv2.contourArea(hull)
        # Moments
        m = cv2.moments(mask, binaryImage=True)
        # Area
        area = m['m00']
        # Solidity
        solidity = area / hull_area if hull_area != 0 else 1
        # Perimeter
        perimeter = cv2.arcLength(obj, closed=True)
        # Bounding rectangle
        x, y, width, height = cv2.boundingRect(obj)
        # Centroid/Center of Mass
        cmx = m['m10'] / m['m00']
        cmy = m['m01'] / m['m00']
        # Aspect ratio
        aspect_ratio = width / float(height) if height != 0 else 0
        # Circularity (4π * area / perimeter^2)
        if perimeter > 0:
            circularity = (4 * np.pi * area) / (perimeter ** 2)
        # Shape Factor (perimeter^2 / (4π * area))
        if area > 0:
            shape_factor = (perimeter ** 2) / (4 * np.pi * area)

        # Hu Moments (seven invariant moments)
        # hu_moments = cv2.HuMoments(m).flatten()

        # Store Hu Moments in the moments dictionary
        # for i, hu_moment in enumerate(hu_moments):
        #     moments[f"hu_moment_{i + 1}"] = hu_moment

    # Store outputs
    outputs.add_metadata(term="image_height", datatype=int, value=np.shape(img)[0])
    outputs.add_metadata(term="image_width", datatype=int, value=np.shape(img)[1])
    outputs.add_observation(sample=label, variable='area', trait='area',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(area, "area"), label=params.unit)
    outputs.add_observation(sample=label, variable='convex_hull_area', trait='convex hull area',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(hull_area, "area"), label=params.unit)
    outputs.add_observation(sample=label, variable='solidity', trait='solidity',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=float,
                            value=solidity, label='none')
    outputs.add_observation(sample=label, variable='perimeter', trait='perimeter',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(perimeter), label=params.unit)
    outputs.add_observation(sample=label, variable='width', trait='width',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(width), label=params.unit)
    outputs.add_observation(sample=label, variable='height', trait='height',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(height), label=params.unit)
    outputs.add_observation(sample=label, variable='longest_path', trait='longest path',
                            method='plantcv.plantcv.analyze.size', scale=params.unit, datatype=int,
                            value=_scale_size(float(longest_path)), label=params.unit)
    outputs.add_observation(sample=label, variable='center_of_mass', trait='center of mass',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=tuple,
                            value=(cmx, cmy), label=("x", "y"))
    outputs.add_observation(sample=label, variable='convex_hull_vertices', trait='convex hull vertices',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=int,
                            value=hull_vertices, label='none')
    outputs.add_observation(sample=label, variable='object_in_frame', trait='object in frame',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=bool,
                            value=in_bounds, label='none')

    # Add new features: Circularity, Shape Factor, Aspect Ratio, and Hu Moments
    outputs.add_observation(sample=label, variable='circularity', trait='circularity',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=float,
                            value=circularity, label='none')
    outputs.add_observation(sample=label, variable='shape_factor', trait='shape_factor',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=float,
                            value=shape_factor, label='none')
    outputs.add_observation(sample=label, variable='aspect_ratio', trait='aspect_ratio',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=float,
                            value=aspect_ratio, label='none')

    # Store Hu Moments from moments dictionary
    # for moment_name, moment_value in moments.items():
    #     outputs.add_observation(sample=label, variable=moment_name, trait=moment_name,
    #                             method='plantcv.plantcv.analyze.size', scale='none', datatype=float,
    #                             value=moment_value, label='none')
    # Store the contour object in outputs
    outputs.add_observation(sample=label, variable='contour', trait='contour',
                            method='plantcv.plantcv.analyze.size', scale='none', datatype=list,
                            value=contour_list, label='none')
    return plt_img


def _longest_axis(height, width, hull, cmx, cmy):
    """
    Calculate the line through center of mass and point on the convex hull that is furthest away

    :param height: int
    :param width: int
    :param hull: numpy.ndarray
    :param cmx: int
    :param cmy: int
    :return caliper_length: int
    """
    background = np.zeros((height, width, 3), np.uint8)
    background1 = np.zeros((height, width), np.uint8)
    background2 = np.zeros((height, width), np.uint8)
    # Longest Axis: line through center of mass and point on the convex hull that is furthest away
    cv2.circle(background, (int(cmx), int(cmy)), 4, (255, 255, 255), -1)
    center_p = cv2.cvtColor(background, cv2.COLOR_BGR2GRAY)
    _, centerp_binary = cv2.threshold(center_p, 0, 255, cv2.THRESH_BINARY)
    centerpoint, _ = _cv2_findcontours(bin_img=centerp_binary)

    dist = []
    vhull = np.vstack(hull)

    for i, c in enumerate(vhull):
        xy = tuple(int(ci) for ci in c)
        pptest = cv2.pointPolygonTest(centerpoint[0], xy, measureDist=True)
        dist.append(pptest)

    abs_dist = np.absolute(dist)
    max_i = np.argmax(abs_dist)

    caliper_max_x, caliper_max_y = list(tuple(vhull[max_i]))
    caliper_mid_x, caliper_mid_y = [int(cmx), int(cmy)]

    xdiff = float(caliper_max_x - caliper_mid_x)
    ydiff = float(caliper_max_y - caliper_mid_y)

    # Set default values
    slope = 1

    if xdiff != 0:
        slope = float(ydiff / xdiff)
    b_line = caliper_mid_y - (slope * caliper_mid_x)

    if slope != 0:
        xintercept = int(-b_line / slope)
        xintercept1 = int((height - b_line) / slope)
        if 0 <= xintercept <= width and 0 <= xintercept1 <= width:
            cv2.line(background1, (xintercept1, height), (xintercept, 0), (255), params.line_thickness)
        elif xintercept < 0 or xintercept > width or xintercept1 < 0 or xintercept1 > width:
            yintercept = int(b_line)
            yintercept1 = int((slope * width) + b_line)
            cv2.line(background1, (0, yintercept), (width, yintercept1), (255), 5)
    else:
        cv2.line(background1, (width, caliper_mid_y), (0, caliper_mid_y), (255), params.line_thickness)

    _, line_binary = cv2.threshold(background1, 0, 255, cv2.THRESH_BINARY)

    cv2.drawContours(background2, [hull], -1, (255), -1)
    _, hullp_binary = cv2.threshold(background2, 0, 255, cv2.THRESH_BINARY)

    caliper = cv2.multiply(line_binary, hullp_binary)

    caliper_y, caliper_x = np.array(caliper.nonzero())
    caliper_matrix = np.vstack((caliper_x, caliper_y))
    caliper_transpose = np.transpose(caliper_matrix)
    caliper_length = len(caliper_transpose)

    caliper_transpose1 = np.lexsort((caliper_y, caliper_x))
    caliper_transpose2 = [(caliper_x[i], caliper_y[i]) for i in caliper_transpose1]
    caliper_transpose = np.array(caliper_transpose2)

    return caliper_length, caliper_transpose


def process_image_glcm(img, nbit=64, mi=0, ma=255, slide_window=3, step=[2], angle=[0]):
    # 如果只有单波段图像，调整为多波段格式
    if len(img.shape) == 2:
        img = np.expand_dims(img, axis=0)

    bands, h, w = img.shape
    print(f"Image has {bands} bands.")

    # 初始化存储所有波段特征的列表
    all_band_features = []

    for b in range(bands):
        print(f"Processing band {b + 1}...")

        # 获取当前波段并归一化到 0-255
        band = img[b]
        band = np.uint8(255.0 * (band - np.min(band)) / (np.max(band) - np.min(band)))

        # 计算 GLCM
        glcm = get_glcm.calcu_glcm(band, mi, ma, nbit, slide_window, step, angle)

        # 计算当前波段的 GLCM 特征
        features = []
        for i in range(glcm.shape[2]):
            for j in range(glcm.shape[3]):
                glcm_cut = glcm[:, :, i, j, :, :]
                features.append(get_glcm.calcu_glcm_mean(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_variance(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_homogeneity(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_contrast(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_dissimilarity(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_entropy(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_energy(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_correlation(glcm_cut, nbit))
                features.append(get_glcm.calcu_glcm_Auto_correlation(glcm_cut, nbit))

        # 堆叠当前波段的所有特征
        features = np.array(features)
        all_band_features.append(features)

    # 将所有波段的特征堆叠
    all_band_features = np.vstack(all_band_features)

    return all_band_features


def process_images(full_bil_path, mask, outdir='None', debug="None"):
    """
    处理光谱图像和掩膜文件，进行分析并保存结果。

    Parameters:
    full_bil_path (str): 光谱图像文件路径，包含 .bil 文件。
    mask (numpy.ndarray): 16位标签掩膜图像，每个对象的像素值对应其标签ID。
    outdir (str): 输出结果文件夹路径。
    debug (str): 调试模式（'plot' 或 'print'，默认为 'print'）。

    Returns:
    combined_data: 分析结果的数据列表
    """

    # 定义选项类
    class options:
        def __init__(self):
            self.image = full_bil_path  # 输入的光谱图像路径
            self.debug = debug  # 设置调试模式：'plot' 或 'print'
            self.writeimg = False  # 是否保存图片
            self.result = outdir  # 输出结果文件路径
            self.outdir = outdir  # 输出文件夹路径

    # 初始化参数
    args = options()

    # 设置 PlantCV 的全局参数
    pcv.params.debug = args.debug
    pcv.params.dpi = 100  # 设置 DPI
    pcv.params.text_size = 1  # 文本大小
    pcv.params.text_thickness = 1  # 文本粗细

    # 读取光谱图像（ENVI 格式）
    spectral_array = pcv.readimage(filename=args.image, mode='envi')
    #过曝区域掩膜
    bath_100 = spectral_array.array_data[:, :, 100]
    bath_over = pcv.threshold.binary(gray_img=bath_100, threshold=11000, object_type='dark')

    # mask已经是16位标签掩膜，直接使用
    # 确保掩膜是正确的数据类型
    if mask.dtype != np.uint16:
        mask = mask.astype(np.uint16)

    # 获取标签数量和标签值
    unique_labels = np.unique(mask)
    # 移除背景标签（通常是0）
    unique_labels = unique_labels[unique_labels != 0]
    num = len(unique_labels)

    print(f"Number of objects in mask: {num}")

    # 转换过曝掩膜为二值格式（0和1）
    bath_over_binary = (bath_over > 0).astype(np.uint16)

    # 创建用于形状分析的标签掩膜（去除过曝区域的版本）
    # labeled_mask = mask * bath_over_binary
    labeled_mask = mask

    # 获取剩余的有效标签
    valid_labels = np.unique(labeled_mask)
    valid_labels = valid_labels[valid_labels != 0]
    num_valid = len(valid_labels)

    print(f"Number of objects after removing overexposed regions: {num_valid}")

    # 分析每个对象的形状和大小
    # 使用二值图像用于可视化（将标签掩膜转换为显示格式）
    binary_img = (mask > 0).astype(np.uint8) * 255
    shape_img = size(img=binary_img, labeled_mask=labeled_mask, n_labels=num_valid)

    # 分析光谱反射率
    # 应用过曝掩膜
    labeled_spectral_mask = mask * bath_over_binary

    spectral_hist = pcv.analyze.spectral_reflectance(
        hsi=spectral_array,
        labeled_mask=labeled_spectral_mask,
        n_labels=num_valid,
        label=None
    )

    observations = pcv.outputs.observations
    # 将结果转换为列表
    combined_data = process_plantcv_outputs(observations)

    return combined_data

# # 示例：批量处理指定路径下的光谱图像和掩膜文件
# bil_path = r'D:\WQ\test\Traindata-05'
# mask_path = r'D:\WQ\test\mask'
# outdir = r"D:\WQ\test"  # 输出文件夹路径
#
# process_images(bil_path, mask_path, outdir)