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 DotaOperator import DotaObj,createDota,readDotaFile,writerDotaFile
import argparse
import math
from math import ceil, floor
import random

##########################################################################
# 函数区
##########################################################################
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 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("ENVI")
    dataset = driver.Create(output_path, im_width, im_height, 1, gdal.GDT_Byte)

    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(".bin",".png")
    Image.fromarray(im_data).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), 1024, 1024,
            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 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: 切片类型
    """
    clusterDict = {} # clusterDict[i]={"p":[sx,sy],"id":[]}
    for  did in range(len(dotalist)):
        centerX, centerY, minX, minY, maxX, maxY=getExternCenter(dotalist, [did])
        # 构建随机平移
        EdgeX=maxX-minX
        EdgeY=maxY-minY
        random_numX = random.random()*2-1
        random_numY = random.random()*2-1

        offsetX=int(math.floor((pitchSize/2-EdgeX)*random_numX)) if pitchSize/2 > EdgeX else 0
        offsetY=int(math.floor((pitchSize/2-EdgeY)*random_numY)) if pitchSize/2 > EdgeY else 0

        offsetX= offsetX if centerX+offsetX+pitchSize/2 < imgwidth else math.ceil(imgwidth-pitchSize/2-centerX-1)
        offsetY= offsetY if centerY+offsetY+pitchSize/2 < imgheight else math.ceil(imgheight-pitchSize/2-centerY-1)

        offsetX= math.floor(pitchSize/2-centerX) if centerX+offsetX-pitchSize/2 < 0 else offsetX
        offsetY= math.floor(pitchSize/2-centerY) if centerY+offsetY-pitchSize/2 < 0 else offsetY

        minX=int(centerX+offsetX-pitchSize/2)
        maxY=int(centerY+offsetY-pitchSize/2)

        clusterDict[did]={"p":[minX,maxY],
                          "id":[did]}

    for cid in clusterDict:
        [minX,minY]=clusterDict[cid]["p"]
        Abox=[minX,minY,pitchSize,pitchSize]
        for did in range(len(dotalist)):
            [centerX, centerY, minX, minY, maxX, maxY] = getExternCenter(dotalist, [did])
            EdgeX = maxX - minX
            EdgeY = maxY - minY
            Bbox=[minX, minY,EdgeX,EdgeY]

            if check_B_in_A(Abox,Bbox):
                if did not in clusterDict[cid]["id"]:
                    clusterDict[cid]["id"].append(did)

    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)

    # 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, colid]
        ax.set_title(str(cid))
        sliceData=tiff_data[sy:(sy+1024),sx:(sx+1024)]
        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"]
        sliceData = tiff_data[sy:(sy + 1024), sx:(sx + 1024)]
        outbinname="{}_{}.bin".format(tiff_basename,str(cid).zfill(3))
        outlabelname="{}_{}.txt".format(tiff_basename,str(cid).zfill(3))
        # 获取样本列表
        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)
    # drawSliceRasterPrivew(tiff_data, dotalist, clusterDict)
    ncols=int(len(clusterDict)/9+1)
    # drawSlictplot(clusterDict, dotalist, tiff_data, 9, ncols)
    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'F:\天仪SAR卫星数据集\舰船数据\bc2-sp-org-vv-20250205t032055-021998-000036-0055ee-01.bin', help='输入tiff的bin文件')
    parser.add_argument('-l', '--labelfilepath',type=str,default=r"F:\天仪SAR卫星数据集\舰船数据\标注\bc2-sp-org-vv-20250205t032055-021998-000036-0055ee-01_LC.txt", help='输入标注')
    parser.add_argument('-o', '--outfolderpath',type=str,default=r'F:\天仪SAR卫星数据集\舰船数据\切片结果', help='切片文件夹地址')
    parser.add_argument('-f', '--filterlabel',type=str,default=r'JLC;MLC', 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)











