import os
import argparse
from osgeo import ogr,gdal
from matplotlib import pyplot as plt
from osgeo import gdal
import matplotlib
import matplotlib.patches as patches
from osgeo import gdal
from PIL import Image
from scipy.spatial import cKDTree
import numpy as np
from tools.DotaOperator import DotaObj,createDota,readDotaFile,writerDotaFile
import argparse
import math
from math import ceil, floor

##########################################################################
# 参数区
##########################################################################
SliceSize=5000

##########################################################################
# 函数区
##########################################################################
def read_tif(path):
    dataset = gdal.Open(path)  # 打开TIF文件
    if dataset is None:
        print("无法打开文件")
        return None, None, None

    cols = dataset.RasterXSize  # 图像宽度
    rows = dataset.RasterYSize  # 图像高度
    bands = dataset.RasterCount
    im_proj = dataset.GetProjection()  # 获取投影信息
    im_Geotrans = dataset.GetGeoTransform()  # 获取仿射变换信息
    im_data = dataset.ReadAsArray(0, 0, cols, rows)  # 读取栅格数据为NumPy数组
    print("行数：", rows)
    print("列数：", cols)
    print("波段：", bands)
    del dataset  # 关闭数据集
    return im_proj, im_Geotrans, im_data


def Strech_SquareRoot(im_data):
    # 判断是否为dB
    # immask = np.isfinite(im_data)
    # imvail_data = im_data[immask]
    # minvalue = np.percentile(imvail_data,30)
    # if minvalue<0 :
    #     im_data=np.power(10.0,im_data/10.0)

    im_data=np.sqrt(im_data)
    immask = np.isfinite(im_data)
    imvail_data = im_data[immask]

    minvalue=np.nanmin(imvail_data)
    maxvalue=np.nanmax(imvail_data)
    minvalue_01Prec = np.percentile(imvail_data, 0.1)  # 20250904 1%拉伸
    maxvalue_999Prec = np.percentile(imvail_data, 99.9)
    print('sqrt root min - max ', minvalue,maxvalue)
    if (maxvalue-minvalue)/(maxvalue_999Prec-minvalue_01Prec)>3: # 表示 拉伸之后，像素值绝大部分很有可能集中在 80
        minvalue=minvalue_01Prec
        maxvalue=maxvalue_999Prec
        print('sqrt root min(0.1) - max(99.9) ', minvalue, maxvalue)


    im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254 + 1
    im_data = np.clip(im_data, 0, 255)

    return im_data.astype(np.uint8)

def write_envi(im_data, im_geotrans, im_proj, output_path):
    """
    将数组数据写入ENVI格式文件
    :param im_data: 输入的numpy数组（2D或3D）
    :param im_geotrans: 仿射变换参数（6元组）
    :param im_proj: 投影信息（WKT字符串）
    :param output_path: 输出文件路径（无需扩展名，会自动生成.dat和.hdr）
    """
    im_bands = 1
    im_height, im_width = im_data.shape
    # 创建ENVI格式驱动
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(output_path, im_width, im_height, 1, gdal.GDT_Float32)

    if dataset is not None:
        dataset.SetGeoTransform(im_geotrans)  # 设置地理变换参数
        dataset.SetProjection(im_proj)  # 设置投影

    dataset.GetRasterBand(1).WriteArray(im_data)

    dataset.FlushCache()  # 确保数据写入磁盘
    dataset = None  # 关闭文件
    outfilepath=output_path.replace(".tiff",".png")
    im_data_uint8=Strech_SquareRoot(im_data)
    Image.fromarray(im_data_uint8).save(outfilepath, compress_level=0)

def geoXY2pixelXY(geo_x, geo_y, inv_gt):
    pixel_x = inv_gt[0] + geo_x * inv_gt[1] + geo_y * inv_gt[2]
    pixel_y = inv_gt[3] + geo_x * inv_gt[4] + geo_y * inv_gt[5]
    return pixel_x, pixel_y

def label2pixelpoints(dotapath,tiff_inv_trans,methodstr,filterlabels):
    dotalist = readDotaFile(dotapath,filterlabels)
    if methodstr=="geolabel":
        for i in range(len(dotalist)):
            geo_x = dotalist[i].x1  # x1
            geo_y = dotalist[i].y1
            pixel_x, pixel_y = geoXY2pixelXY(geo_x, geo_y, tiff_inv_trans)
            dotalist[i].x1 = pixel_x
            dotalist[i].y1 = pixel_y

            geo_x = dotalist[i].x2  # x2
            geo_y = dotalist[i].y2
            pixel_x, pixel_y = geoXY2pixelXY(geo_x, geo_y, tiff_inv_trans)
            dotalist[i].x2 = pixel_x
            dotalist[i].y2 = pixel_y

            geo_x = dotalist[i].x3  # x3
            geo_y = dotalist[i].y3
            pixel_x, pixel_y = geoXY2pixelXY(geo_x, geo_y, tiff_inv_trans)
            dotalist[i].x3 = pixel_x
            dotalist[i].y3 = pixel_y

            geo_x = dotalist[i].x4  # x4
            geo_y = dotalist[i].y4
            pixel_x, pixel_y = geoXY2pixelXY(geo_x, geo_y, tiff_inv_trans)
            dotalist[i].x4 = pixel_x
            dotalist[i].y4 = pixel_y

    print("点数:", len(dotalist))
    return dotalist

def getMaxEdge(dotalist, ids):
    cornpoint = np.zeros((len(ids) * 4, 2))
    for idx in range(len(ids)):
        cornpoint[idx * 4 + 0, 0] = dotalist[ids[idx]].x1
        cornpoint[idx * 4 + 1, 0] = dotalist[ids[idx]].x2
        cornpoint[idx * 4 + 2, 0] = dotalist[ids[idx]].x3
        cornpoint[idx * 4 + 3, 0] = dotalist[ids[idx]].x4

        cornpoint[idx * 4 + 0, 1] = dotalist[ids[idx]].y1
        cornpoint[idx * 4 + 1, 1] = dotalist[ids[idx]].y2
        cornpoint[idx * 4 + 2, 1] = dotalist[ids[idx]].y3
        cornpoint[idx * 4 + 3, 1] = dotalist[ids[idx]].y4

    xedge = np.max(cornpoint[:, 0]) - np.min(cornpoint[:, 0])
    yedge = np.max(cornpoint[:, 1]) - np.min(cornpoint[:, 1])

    edgelen = xedge if xedge > yedge else yedge
    return edgelen

def getExternCenter(dotalist, ids):
    cornpoint = np.zeros((len(ids) * 4, 2))
    for idx in range(len(ids)):
        cornpoint[idx * 4 + 0, 0] = dotalist[ids[idx]].x1
        cornpoint[idx * 4 + 1, 0] = dotalist[ids[idx]].x2
        cornpoint[idx * 4 + 2, 0] = dotalist[ids[idx]].x3
        cornpoint[idx * 4 + 3, 0] = dotalist[ids[idx]].x4

        cornpoint[idx * 4 + 0, 1] = dotalist[ids[idx]].y1
        cornpoint[idx * 4 + 1, 1] = dotalist[ids[idx]].y2
        cornpoint[idx * 4 + 2, 1] = dotalist[ids[idx]].y3
        cornpoint[idx * 4 + 3, 1] = dotalist[ids[idx]].y4

    minX = np.min(cornpoint[:, 0])
    minY = np.min(cornpoint[:, 1])
    maxX = np.max(cornpoint[:, 0])
    maxY = np.max(cornpoint[:, 1])
    centerX = (minX + maxX) / 2
    centerY = (minY + maxY) / 2
    return [centerX, centerY, minX, minY, maxX, maxY]

def drawSliceRasterPrivew(tiff_data,dotalist,clusterDict):
    # 绘制图形
    # 创建图形和坐标轴
    fig, ax = plt.subplots(figsize=(20, 16))
    ax.imshow(tiff_data, cmap='gray')
    # 绘制每个目标的矩形框并标注坐标
    for i in range(len(dotalist)):
        # 提取x和y坐标
        x_coords = [dotalist[i].x1, dotalist[i].x2, dotalist[i].x3, dotalist[i].x4]
        y_coords = [dotalist[i].y1, dotalist[i].y2, dotalist[i].y3, dotalist[i].y4]

        # 计算最小外接矩形（AABB）
        x_min, x_max = min(x_coords), max(x_coords)
        y_min, y_max = min(y_coords), max(y_coords)
        width = x_max - x_min
        height = y_max - y_min

        # 绘制无填充矩形框（仅红色边框）
        rect = patches.Rectangle(
            (x_min, y_min), width, height,
            linewidth=2, edgecolor='red', facecolor='none'  # 关键：facecolor='none'
        )
        ax.add_patch(rect)

        # ax.annotate(f'({x},{y})', xy=(x, y), xytext=(5, 5),
        #                textcoords='offset points', fontsize=10,
        #                bbox=dict(boxstyle='round,pad=0.5', fc='white', alpha=0.8))

        # 在矩形中心标注目标编号
        center_x = sum(x_coords) / 4
        center_y = sum(y_coords) / 4
        ax.text(center_x, center_y, str(i),
                ha='center', va='center', fontsize=6, color='red')

    # 以类别中心为中心绘制四边形
    for k in clusterDict:
        # 绘制无填充矩形框（仅红色边框）
        minX = clusterDict[k]["p"][0]
        minY = clusterDict[k]["p"][1]
        rect = patches.Rectangle(
            (minX , minY), SliceSize, SliceSize,
            linewidth=2, edgecolor='green', facecolor='none'  # 关键：facecolor='none'
        )
        ax.add_patch(rect)
        ax.text(minX+512, minY+512, str(k),
                ha='center', va='center', fontsize=6, color='green')

    plt.tight_layout()
    plt.show()

    print("绘图结束")
    return None

def find_optimal_slices(H, W, boxes, patch_size=1024, max_overlap_rate=0.2):
    """
    Compute optimal slice positions for the image to maximize the number of fully contained rectangular patches (boxes),
    while ensuring the overlap rate between any two slices does not exceed the specified maximum.

    Parameters:
    - H: Height of the image.
    - W: Width of the image.
    - boxes: List of tuples or lists, each containing (x1, y1, x2, y2) where (x1, y1) is the top-left and (x2, y2) is the bottom-right of a rectangular patch.
    - patch_size: Size of each slice (square, e.g., 1024).
    - max_overlap_rate: Maximum allowed overlap rate (e.g., 0.2).

    Returns:
    - slices: List of (sx, sy) starting positions for the slices.
    - covered_count: Number of patches that appear fully in at least one slice.
    """
    overlap_max = patch_size * max_overlap_rate
    stride = patch_size - floor(overlap_max)  # Ensures overlap <= max_overlap_rate in linear dimensions

    N = len(boxes)
    x_covered = [set() for _ in range(stride)]
    y_covered = [set() for _ in range(stride)]

    for i in range(N):
        x1, y1, x2, y2 = boxes[i]
        b_w = x2 - x1
        b_h = y2 - y1

        # For x-dimension
        lx = max(0, x2 - patch_size)
        rx = min(W - patch_size, x1)
        if lx <= rx:
            start_x = ceil(lx)
            end_x = floor(rx)
            l_x = end_x - start_x + 1
            if l_x > 0:
                if l_x >= stride:
                    for ox in range(stride):
                        x_covered[ox].add(i)
                else:
                    for sx in range(start_x, end_x + 1):
                        ox = sx % stride
                        x_covered[ox].add(i)

        # For y-dimension
        ly = max(0, y2 - patch_size)
        ry = min(H - patch_size, y1)
        if ly <= ry:
            start_y = ceil(ly)
            end_y = floor(ry)
            l_y = end_y - start_y + 1
            if l_y > 0:
                if l_y >= stride:
                    for oy in range(stride):
                        y_covered[oy].add(i)
                else:
                    for sy in range(start_y, end_y + 1):
                        oy = sy % stride
                        y_covered[oy].add(i)

    # Find the best offset pair (ox, oy) that maximizes covered patches
    max_covered = 0
    best_ox = 0
    best_oy = 0
    for ox in range(stride):
        for oy in range(stride):
            current_covered = len(x_covered[ox] & y_covered[oy])
            if current_covered > max_covered:
                max_covered = current_covered
                best_ox = ox
                best_oy = oy

    # Generate the slice positions using the best offsets and stride
    slices = []
    sx = best_ox
    while sx + patch_size <= W:
        sy = best_oy
        while sy + patch_size <= H:
            slices.append((sx, sy))
            sy += stride
        sx += stride

    return slices, max_covered

def check_B_in_A(A,B):
    """
    判断A包含B
    :param A: [x0,y0.w.h]
    :param B: [x0,y0.w.h]
    :return:
    """
    # 解构矩形A和B的参数
    Ax0, Ay0, Aw, Ah = A
    Bx0, By0, Bw, Bh = B

    # 计算矩形A和B的右边界和下边界
    Ax1 = Ax0 + Aw
    Ay1 = Ay0 + Ah

    Bx1 = Bx0 + Bw
    By1 = By0 + Bh

    # 判断B是否完全在A内部
    return (Bx0 >= Ax0) and (Bx1 <= Ax1) and (By0 >= Ay0) and (By1 <= Ay1)


##########################################################################
# 切分算法流程图
##########################################################################

def getclusterDict(dotalist,imgheight,imgwidth,pitchSize=1024,max_overlap_rate=0.2):
    """
    生成切片数据
    :param dotalist: 样本集
    :param imgheight: 图像高度
    :param imgwidth: 图像宽度
    :return: 切片类型
    """
    boxs=[]
    for i in range(len(dotalist)):
        xs=np.array([dotalist[i].x1,dotalist[i].x2,dotalist[i].x3, dotalist[i].x4])
        ys=np.array([dotalist[i].y1,dotalist[i].y2,dotalist[i].y3, dotalist[i].y4])
        x1=np.min(xs)
        x2=np.max(xs)
        y1=np.min(ys)
        y2=np.max(ys)
        boxs.append([x1,y1,x2,y2]) # x1, y1, x2, y2 = boxes[i]

    slices, max_covered=find_optimal_slices(imgheight,imgwidth,boxs,pitchSize,max_overlap_rate)

    clusterDict={}

    waitContaindota=[]
    hasContainIds=[]
    for i in range(len(slices)):
        sx,sy=slices[i]
        clusterDict[i]={"p":[sx,sy],"id":[]}
        slicesExten=[sx,sy,SliceSize,SliceSize]
        for ids in range(len(dotalist)):
            if ids in hasContainIds:
                continue
            else:
                [centerX, centerY, minX, minY, maxX, maxY]=getExternCenter(dotalist, [ids])
                dotaExtend=[minX,minY,maxX-minX,maxY-minY]
                if check_B_in_A(slicesExten,dotaExtend):
                    print("True: ", slicesExten, dotaExtend)
                    clusterDict[i]["id"].append(ids)
                    # hasContainIds.append(ids)
                else:
                    print("False: ",slicesExten,dotaExtend)

    for ids in range(len(dotalist)):
        if ids in hasContainIds:
            continue
        else:
            waitContaindota.append(ids)
            print("No in slice dota : ",str(dotalist[ids]) )

    print("no process ids ",str(waitContaindota))
    return clusterDict


def drawSlictplot(clusterDict,dotalist,tiff_data,nrows=10,ncols=9):
    """
    :param clusterDict:  clusterDict[i]={"p":[sx,sy],"id":[]}
    :param dotalist:  （x1, y1, x2, y2, x3, y3, x4, y4 clsname diffcule）
    :return:
    """
    fig, axes  = plt.subplots(nrows=nrows,ncols=ncols,figsize=(20, 16))
    plt.tight_layout(pad=3.0)
    print(tiff_data.shape)
    # 9*10
    subid=0
    for cid in clusterDict:
        sx,sy=clusterDict[cid]["p"]
        colid=subid//nrows
        rowid=subid%nrows
        subid=subid+1
        ax = axes[rowid]
        ax.set_title(str(cid))
        sliceData=tiff_data[sy:(sy+SliceSize),sx:(sx+SliceSize)]
        ax.imshow(sliceData, cmap='gray')

        for did in clusterDict[cid]["id"] :
            # 提取x和y坐标
            x_coords = [dotalist[did].x1-sx, dotalist[did].x2-sx, dotalist[did].x3-sx, dotalist[did].x4-sx]
            y_coords = [dotalist[did].y1-sy, dotalist[did].y2-sy, dotalist[did].y3-sy, dotalist[did].y4-sy]

            # 计算最小外接矩形（AABB）
            x_min, x_max = min(x_coords), max(x_coords)
            y_min, y_max = min(y_coords), max(y_coords)
            width = x_max - x_min
            height = y_max - y_min


            # 绘制无填充矩形框（仅红色边框）
            rect = patches.Rectangle(
                (x_min, y_min), width, height,
                linewidth=2, edgecolor='red', facecolor='none'  # 关键：facecolor='none'
            )
            ax.add_patch(rect)

            # 在矩形中心标注目标编号
            center_x = x_min+width/2
            center_y = y_min+height/2
            ax.text(center_x, center_y, str(did),
                    ha='center', va='center', fontsize=6, color='red')


    plt.tight_layout()
    plt.show()

    print("绘图结束")
    return None


def slictDataAndOutlabel(clusterDict,dotalist,tiff_data,tiff_basename,outfolderpath,im_geotrans, im_proj):
    """
    切分标签，输出结果与文件
    :param clusterDict:
    :param dotalist:
    :param tiff_data:
    :param tiff_name:
    :param outfolderpath:
    :return:
    """
    for cid in clusterDict:
        sx, sy = clusterDict[cid]["p"]
        if len(clusterDict[cid]["id"])==0:
            continue
        sliceData = tiff_data[sy:(sy + SliceSize), sx:(sx + SliceSize)]
        outbinname="{}_{}.tiff".format(tiff_basename,cid)
        outlabelname="{}_{}.txt".format(tiff_basename,cid)
        # 获取样本列表
        outdotalist=[]
        for did in clusterDict[cid]["id"] :
            x1=dotalist[did].x1-sx
            x2=dotalist[did].x2-sx
            x3=dotalist[did].x3-sx
            x4=dotalist[did].x4-sx
            y1=dotalist[did].y1-sy
            y2=dotalist[did].y2-sy
            y3=dotalist[did].y3-sy
            y4=dotalist[did].y4-sy
            tempdota=createDota(x1,y1,x2,y2,x3,y3,x4,y4,dotalist[did].clsname,dotalist[did].difficulty)
            outdotalist.append(tempdota)

        outlabelpath=os.path.join(outfolderpath,outlabelname)
        outbinpath=os.path.join(outfolderpath,outbinname)

        temp_im_geotrans=[tempi for tempi in im_geotrans]
        # 处理 0,3
        temp_im_geotrans[0]=im_geotrans[0]+sx*im_geotrans[1]+im_geotrans[2]*sy # x
        temp_im_geotrans[3]=im_geotrans[3]+sx*im_geotrans[4]+im_geotrans[5]*sy # y
        write_envi(sliceData,temp_im_geotrans,im_proj,outbinpath)
        writerDotaFile(outdotalist,outlabelpath)

##########################################################################
# 处理流程图
##########################################################################
def DataSampleSliceRasterProcess(inbinfile,labelfilepath,outfolderpath,methodstr,filterlabels):
    tiff_proj, tiff_trans, tiff_data = read_tif(inbinfile)
    tiff_inv_trans = gdal.InvGeoTransform(tiff_trans)
    dotalist=label2pixelpoints(labelfilepath,tiff_inv_trans,methodstr,filterlabels)
    imgheight, imgwidth=tiff_data.shape
    clusterDict=getclusterDict(dotalist,imgheight,imgwidth,SliceSize,0.25)
    # drawSliceRasterPrivew(tiff_data, dotalist, clusterDict)
    nrows=int(len(clusterDict)/1+1)
    # drawSlictplot(clusterDict, dotalist, tiff_data, nrows, 1)
    tiff_name=os.path.basename(inbinfile)
    tiff_basename=os.path.splitext(tiff_name)[0]
    slictDataAndOutlabel(clusterDict, dotalist, tiff_data, tiff_basename, outfolderpath, tiff_trans, tiff_proj)

    pass



def getParams():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--inbinfile',type=str,default=r'D:\港口\Geo_bc2-sm-org-vv-20231016t135315-008424-0020e8-01.tiff', help='输入tiff的bin文件')
    parser.add_argument('-l', '--labelfilepath',type=str,default=r"D:\港口\港口dota\Geo_bc2-sm-org-vv-20231016t135315-008424-0020e8-01.military_harbor.txt", help='输入标注')
    parser.add_argument('-o', '--outfolderpath',type=str,default=r'D:\港口\切片结果', help='切片文件夹地址')
    parser.add_argument('-f', '--filterlabel',type=str,default=r'mix_airport;civil_harbor;military_harbor;no_harbor', help='标签过滤')
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        '--geolabel',
        action='store_const',
        const='geolabel',
        dest='method',
        help='标注坐标点为地理坐标'
    )
    group.add_argument(
        '--pixellabel',
        action='store_const',
        const='pixellabel',
        dest='method',
        help='标注坐标系统为输入影像的像空间坐标'
    )
    parser.set_defaults(method='geolabel')


    args = parser.parse_args()
    return args

if __name__ == '__main__':
    try:
        parser = getParams()
        inbinfile=parser.inbinfile
        labelfilepath=parser.labelfilepath
        outfolderpath=parser.outfolderpath
        methodstr= parser.method
        filterlabels=parser.filterlabel.strip().split(';')
        print('inbinfile=',inbinfile)
        print('labelfilepath=',labelfilepath)
        print('outfolderpath=',outfolderpath)
        print('methodstr=',methodstr)
        print('filterlabels=',filterlabels)
        DataSampleSliceRasterProcess(inbinfile, labelfilepath, outfolderpath,methodstr,filterlabels)
        print("样本切分完成")
        exit(2)
    except Exception as e:
        print(e)
        exit(3)












