"""
#!/usr/bin/python3.9
# -*- coding: utf-8 -*-
@Project :pythonCodeSnippet
@File    :aiTools.py
@IDE     :PyCharm
@Author  :chenxw
@Date    :2023/9/13 11:53
@Descr:
"""
import os

import cv2
import numpy as np
from PIL.Image import Image
from osgeo import ogr, gdal
from vgis_gis.shpTools import ShpFileOperator
from vgis_rs.rsTools import RsToolsOperatoer

from vgis_utils.vgis_image.imageTools import ImageHelper
from vgis_utils.vgis_list.listTools import ListHelper


class AIHelper:
    def __init__(self):
        pass

    @staticmethod
    # 计算一个 ground truth 边界盒和 k 个先验框(Anchor)的交并比(IOU)值。
    def __iou(box, clusters):
        """
        计算一个 ground truth 边界盒和 k 个先验框(Anchor)的交并比(IOU)值。
        参数box: 元组或者数据，代表 ground truth 的长宽。
        参数clusters: 形如(k,2)的numpy数组，其中k是聚类Anchor框的个数
        返回：ground truth和每个Anchor框的交并比。
        """
        x = np.minimum(clusters[:, 0], box[0])
        y = np.minimum(clusters[:, 1], box[1])
        if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
            raise ValueError("Box has no area")
        intersection = x * y
        box_area = box[0] * box[1]
        cluster_area = clusters[:, 0] * clusters[:, 1]
        iou_ = intersection / (box_area + cluster_area - intersection)
        return iou_

    @staticmethod
    # 计算一个ground truth和k个Anchor的交并比的均值。
    def avg_iou(boxes, clusters):
        """
        计算一个ground truth和k个Anchor的交并比的均值。
        """
        return np.mean([np.max(AIHelper.__iou(boxes[i], clusters)) for i in range(boxes.shape[0])])

    @staticmethod
    # 利用IOU值进行K-means聚类
    def kmeans(boxes, k, dist=np.median):
        """
        利用IOU值进行K-means聚类
        参数boxes: 形状为(r, 2)的ground truth框，其中r是ground truth的个数
        参数k: Anchor的个数
        参数dist: 距离函数
        返回值：形状为(k, 2)的k个Anchor框
        """
        # 即是上面提到的r
        rows = boxes.shape[0]
        # 距离数组，计算每个ground truth和k个Anchor的距离
        distances = np.empty((rows, k))
        # 上一次每个ground truth"距离"最近的Anchor索引
        last_clusters = np.zeros((rows,))
        # 设置随机数种子
        np.random.seed()

        # 初始化聚类中心，k个簇，从r个ground truth随机选k个
        clusters = boxes[np.random.choice(rows, k, replace=False)]
        # 开始聚类
        while True:
            # 计算每个ground truth和k个Anchor的距离，用1-IOU(box,anchor)来计算
            for row in range(rows):
                distances[row] = 1 - AIHelper.__iou(boxes[row], clusters)
            # 对每个ground truth，选取距离最小的那个Anchor，并存下索引
            nearest_clusters = np.argmin(distances, axis=1)
            # 如果当前每个ground truth"距离"最近的Anchor索引和上一次一样，聚类结束
            if (last_clusters == nearest_clusters).all():
                break
            # 更新簇中心为簇里面所有的ground truth框的均值
            for cluster in range(k):
                clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)
            # 更新每个ground truth"距离"最近的Anchor索引
            last_clusters = nearest_clusters

        return clusters

    @staticmethod
    # 获取两个矩形框的重叠面积
    #     rect1 = [0, 0, 5, 5]
    #     rect2 = [3, 3, 8, 8]
    def get_overlay_radio_of_two_box(rect1, rect2):
        x_overlap = max(0, min(rect1[2], rect2[2]) - max(rect1[0], rect2[0]))
        y_overlap = max(0, min(rect1[3], rect2[3]) - max(rect1[1], rect2[1]))
        overlap_area = x_overlap * y_overlap
        rect1_eara = (rect1[2] - rect1[0]) * (rect1[3] - rect1[1])
        overlap_radio = overlap_area / rect1_eara
        return overlap_radio

    @staticmethod
    # yolov8结果框转化为shapefile文件
    # xyxy_box_list:[[pixel_minx,pixel_miny,pixel_maxx,pixel_maxy],[pixel_minx2,pixel_miny2,pixel_maxx2,pixel_maxy2]]
    # image_shape:(width,height)
    # image_map_env:(map_minx,map_miny,map_maxx,map_maxy)
    def convert_yolo_box_to_shape(yolov_xyxy_box_list, names, image_path, image_map_env, map_epsg, out_shp_path):
        img = cv2.imread(image_path)
        image_shape = img.shape

        map_coords_list = []
        for xyxy_box in yolov_xyxy_box_list:
            pixel_x_min = xyxy_box[0]
            pixel_y_min = xyxy_box[1]
            pixel_x_max = xyxy_box[2]
            pixel_y_max = xyxy_box[3]
            conf = xyxy_box[4]
            cls = xyxy_box[5]
            name = names[cls]
            image_height = image_shape[0]
            image_width = image_shape[1]
            image_minx_map = image_map_env[0]
            image_miny_map = image_map_env[1]
            image_maxx_map = image_map_env[2]
            image_maxy_map = image_map_env[3]
            map_x_min = image_minx_map + (image_maxx_map - image_minx_map) / image_width * pixel_x_min
            map_y_min = image_maxy_map - (image_maxy_map - image_miny_map) / image_height * pixel_y_max
            map_x_max = image_minx_map + (image_maxx_map - image_minx_map) / image_width * pixel_x_max
            map_y_max = image_maxy_map - (image_maxy_map - image_miny_map) / image_height * pixel_y_min
            map_coords_list.append([map_x_min, map_y_min, map_x_max, map_y_max, conf, cls, name])
        ShpFileOperator.convert_box_points_and_attr_into_shp(map_coords_list, map_epsg, out_shp_path)

    @staticmethod
    # mask二值图转shp
    # 通过gdal的栅格转矢量的方法
    def convert_mask_to_shape(mask_image_path: str, mask_image_tranform: list, file_prefix: str, shp_file_path: str,
                              logger: object) -> str:
        """
        将mask二值图转换为shp文件，通过GDAL的栅格转矢量方法

        :param mask_image_path: 二值图栅格文件路径.
        :param mask_image_tranform: 栅格文件的几何转换六参数，格式为[左上角点坐标X, x方向的分辨, 旋转系数:如果为0就是标准的正北向图像, 左上角点坐标Y, 旋转系数:如果为0就是标准的正北向图像, Y方向的分辨率].
        :param file_prefix: 文件前缀
        :param shp_file_path: 转换的shp文件路径
        :param logger: 日志对象
        :return: 转后的矢量shp路径
        """
        if logger is not None:
            logger.info("对二值图转换为shp文件")
        shp_file_dir = os.path.join(os.getcwd(), "shp", "convert_by_mask")
        if not os.path.exists(shp_file_dir):
            os.makedirs(shp_file_dir)
        (file_pre_path, temp_filename) = os.path.split(mask_image_path)
        (shot_name, file_ext) = os.path.splitext(temp_filename)
        if shp_file_path is None:
            shp_file_path = os.path.join(shp_file_dir, '{}.shp'.format(shot_name))
        if logger is not None:
            logger.info("shp文件路径：{}".format(shp_file_path))
        RsToolsOperatoer.raster_to_vector(mask_image_path, mask_image_tranform, shp_file_path, "ESRI Shapefile")
        return shp_file_path

    @staticmethod
    # mask二值图转polygon
    # 通过opencv的轮廓检测（等值线捕获）方法
    # 返回[[[x1,y1],[x2,y2]...[xn,yn]],[...]]
    def mask_to_polygon(mask_file, epsilon_factor=0.001):
        polygon_points = []
        # binary_mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)
        binary_mask = cv2.imdecode(np.fromfile(mask_file, dtype=np.uint8), 0)
        # 二值图边缘线，只能读单波段的灰度图
        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            epsilon = epsilon_factor * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            if len(approx) < 5:
                continue
            each_poly_points = []
            for point in approx:
                x, y = point[0].tolist()
                each_poly_points.append([x, y])
            polygon_points.append(each_poly_points)
        return polygon_points

    @staticmethod
    # mask二值图转polygon
    # 通过opencv的轮廓检测（等值线捕获）方法
    # 提高轮廓检测的精度，你可以使用更高级的图像分割技术，如分水岭分割或GrabCut算法
    # 返回[[[x1,y1],[x2,y2]...[xn,yn]],[...]]
    def mask_to_polygon2(mask_file, epsilon_factor=0.001):

        # 定义GrabCut函数
        def grabcut_segmentation(image):
            # 创建掩码
            mask = np.zeros(image.shape[:2], np.uint8)
            # 创建背景和前景模型
            bgdModel = np.zeros((1, 65), np.float64)
            fgdModel = np.zeros((1, 65), np.float64)
            # 定义矩形区域（这里使用整个图像）
            rect = (0, 0, image.shape[1], image.shape[0])
            # 运行GrabCut算法
            cv2.grabCut(image, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
            # 提取前景
            mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
            return mask2

        polygon_points = []
        binary_mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)

        segmented_mask = grabcut_segmentation(binary_mask)
        # 使用轮廓检测来提取多边形坐标
        contours, _ = cv2.findContours(segmented_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            epsilon = epsilon_factor * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            if len(approx) < 5:
                continue
            each_poly_points = []
            for point in approx:
                x, y = point[0].tolist()
                each_poly_points.append([x, y])
            polygon_points.append(each_poly_points)
        return polygon_points

    @staticmethod
    # 多边形转换为mask二值图
    # polygons:[[(x1,y1),(x2,y2)...(xn,yn)],[...]]
    def polygon_to_mask(img_file, mask_file, polygons):
        def get_image_size(image_file):
            with Image.open(image_file) as img:
                width, height = img.size
                return width, height

        image_width, image_height = get_image_size(img_file)
        image_shape = (image_height, image_width)
        binary_mask = np.zeros(image_shape, dtype=np.uint8)
        for polygon_points in polygons:
            np_polygon = np.array(polygon_points, np.int32)
            np_polygon = np_polygon.reshape((-1, 1, 2))
            cv2.fillPoly(binary_mask, [np_polygon], color=255)
        cv2.imwrite(mask_file, binary_mask)

    @staticmethod
    # 多边形转为mask二值图方法2
    # polygons:['x1,y1,x2,y2...xn,yn','x1,y1,x2,y2...xn,yn'...]
    def polygon_to_mask2(img_file, mask_file, polygons):
        ImageHelper.build_binary_image(img_file, mask_file, polygons)

    @staticmethod
    # 处理AI提取结果，删除多余的提取框
    # 非极大值抑制(Non-maximum supression)简称NMS,其作用是去除冗余的检测框,去冗余手段是剔除与极大值重叠较多的检测框结果
    def handle_ai_result_of_nms(input_shp_path):
        # 注册所有驱动
        gdal.AllRegister()
        # 解决中文路径乱码问题
        gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "NO")
        driver = ogr.GetDriverByName('ESRI Shapefile')
        origin_pFeatureDataset = driver.Open(input_shp_path, 1)
        origin_pFeaturelayer = origin_pFeatureDataset.GetLayer(0)
        copy_pFeatureDataset = ShpFileOperator.get_clone_shape(input_shp_path)
        copy_pFeaturelayer = copy_pFeatureDataset.GetLayer(0)
        copy_featureLayerNum = copy_pFeaturelayer.GetFeatureCount(0)
        # 获取要素
        for t in range(0, copy_featureLayerNum):
            copy_ofeature = copy_pFeaturelayer.GetFeature(t)
            copy_geom = copy_ofeature.GetGeometryRef()
            # copy_geom_wkt = copy_geom.ExportToWkt()
            del_id_list = AIHelper.__find_overlap_feature(copy_geom, origin_pFeaturelayer, 0.5)
        del_id_list_only = ListHelper.remove_same_item(del_id_list)
        for del_id in del_id_list_only:
            print("----正在删除要素ID:" + str(del_id))
            origin_pFeaturelayer.DeleteFeature(del_id)
        strSQL = "REPACK " + str(origin_pFeaturelayer.GetName())
        origin_pFeatureDataset.ExecuteSQL(strSQL, None, "")
        origin_pFeaturelayer = None
        origin_pFeatureDataset = None

    @staticmethod
    # 删除发现的重叠要素
    # 指阈值以上的重叠区域的两个要素，删除面积小的那个要素
    def __find_overlap_feature(search_geom, originFeaturelayer, threshold=0.5):
        del_id_list = []
        origin_feature_num = originFeaturelayer.GetFeatureCount(0)
        search_geom_id = 0
        origin_geom_wkt_list = []
        origin_geom_id_list = []

        for t in range(0, origin_feature_num):
            originFeature = originFeaturelayer.GetFeature(t)
            origin_geom_id = originFeature.GetFID()
            origin_geom = originFeature.GetGeometryRef()
            origin_geom_wkt = origin_geom.ExportToWkt()
            origin_geom_wkt_list.append(origin_geom_wkt)
            origin_geom_id_list.append(origin_geom_id)
            if search_geom.Equal(origin_geom):
                search_geom_id = origin_geom_id

        print("当前查询要素{}".format(str(search_geom_id)))
        if search_geom_id in del_id_list:
            print("----要素{}已经在删除序列中，不再进行重叠查询".format(str(search_geom_id)))
        else:
            for t in range(0, len(origin_geom_wkt_list)):
                origin_geom_id = origin_geom_id_list[t]
                origin_geom = ogr.CreateGeometryFromWkt(origin_geom_wkt_list[t])
                if search_geom_id not in del_id_list:
                    if search_geom_id == origin_geom_id:
                        continue
                    elif search_geom.Contains(origin_geom):
                        if origin_geom_id not in del_id_list:
                            del_id_list.append(origin_geom_id)
                            print("----要素{}包含要素{},删除要素{}".format(str(search_geom_id), str(origin_geom_id),
                                                                           str(origin_geom_id)))
                    elif search_geom.Within(origin_geom):
                        if search_geom_id not in del_id_list:
                            del_id_list.append(search_geom_id)
                            print("----要素{}包含要素{},删除要素{}".format(str(origin_geom_id), str(search_geom_id),
                                                                           str(search_geom_id)))
                    elif search_geom.Intersect(origin_geom):
                        inter_geom = search_geom.Intersection(origin_geom)
                        if ShpFileOperator.get_area_envelop(inter_geom) / ShpFileOperator.get_area_envelop(
                                search_geom) > threshold or ShpFileOperator.get_area_envelop(
                            inter_geom) / ShpFileOperator.get_area_envelop(origin_geom) > threshold:
                            if ShpFileOperator.get_area_envelop(search_geom) > ShpFileOperator.get_area_envelop(
                                    origin_geom):
                                if origin_geom_id not in del_id_list:
                                    del_id_list.append(origin_geom_id)
                                    print("----要素{}与要素{}大面积相交,删除要素{}".format(str(search_geom_id),
                                                                                           str(origin_geom_id),
                                                                                           str(origin_geom_id)))
                            else:
                                if search_geom_id not in del_id_list:
                                    del_id_list.append(search_geom_id)
                                    print("----要素{}与要素{}大面积相交,删除要素{}".format(str(search_geom_id),
                                                                                           str(origin_geom_id),
                                                                                           str(search_geom_id)))
        return del_id_list


if __name__ == '__main__':
    box_points_list = [[660.2479858398438, 280.0595397949219, 701.6737060546875, 320.72589111328125],
                       [559.2014770507812, 221.1412353515625, 592.9896240234375, 254.13226318359375],
                       [477.0056457519531, 309.0809020996094, 505.49664306640625, 337.82135009765625],
                       [517.1990966796875, 194.45872497558594, 551.4634399414062, 227.95953369140625],
                       [721.5293579101562, 321.0753173828125, 757.498291015625, 354.52435302734375],
                       [546.3658447265625, 345.80072021484375, 577.3338012695312, 374.9428405761719],
                       [340.64898681640625, 68.33817291259766, 362.81884765625, 90.0564956665039],
                       [776.2349853515625, 355.7030029296875, 812.4939575195312, 390.0914306640625],
                       [617.8046264648438, 401.6883850097656, 648.726318359375, 430.9305419921875],
                       [682.6495361328125, 443.026123046875, 716.3787231445312, 474.0211486816406],
                       [470.60125732421875, 164.7460479736328, 504.5323181152344, 197.5280303955078],
                       [119.45063018798828, 467.18597412109375, 141.20028686523438, 488.65411376953125],
                       [594.6969604492188, 420.94134521484375, 607.457275390625, 433.63287353515625],
                       [247.9298095703125, 497.092041015625, 259.27301025390625, 508.6060485839844]]
    image_path = 'G:\\AI\\test_data\\storagetank\\storagetank1.jpeg'
    jgw_path = 'G:\\AI\\test_data\\storagetank\\storagetank1.jgw'
    shp_path = 'G:\\AI\\test_data\\storagetank\\storagetank1.shp'

    image_env = RsToolsOperatoer.get_env_of_pic(image_path, jgw_path)
    AIHelper.convert_yolo_box_to_shape(box_points_list, image_path, image_env, 4326, shp_path)
