import rasterio
from rasterio.windows import Window
import geopandas as gpd
from shapely.geometry import box
import numpy as np
import os

# 参数设置
image_path = 'path/to/your/image.tif'  # 替换为您的影像路径（20000x20000）
shapefile_path = 'path/to/your/shapefile.shp'  # 替换为您的图斑Shapefile路径
output_dir = 'path/to/output/tiles'  # 输出切片目录
tile_size = 1024  # 固定切片大小
max_overlap_ratio = 0.2  # 重叠小于20%
min_step = int(tile_size * (1 - max_overlap_ratio))  # 最小步长 ≈820

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 读取影像
with rasterio.open(image_path) as src:
    image_width, image_height = src.width, src.height
    assert image_width == 20000 and image_height == 20000, "影像大小不匹配"
    transform = src.transform  # 地理变换矩阵

    # 读取图斑数据（假设CRS一致）
    gdf = gpd.read_file(shapefile_path)
    # 确保图斑CRS与影像一致，如果不一致需转换：gdf = gdf.to_crs(src.crs)

    # 初始化切片列表，避免重叠过多
    generated_windows = []

    for idx, row in gdf.iterrows():
        geom = row.geometry  # 获取单个图斑多边形

        # 计算图斑的边界框 (minx, miny, maxx, maxy) - 像素坐标
        minx, miny, maxx, maxy = geom.bounds
        bbox_width = maxx - minx
        bbox_height = maxy - miny

        # 如果图斑太大，无法放入单个1024x1024切片，需处理（这里简单跳过，您可修改）
        if bbox_width > tile_size or bbox_height > tile_size:
            print(f"图斑 {idx} 太大 ({bbox_width}x{bbox_height})，跳过或需多切片处理")
            continue

        # 计算覆盖完整图斑的最小窗口：居中扩展到1024x1024
        center_x = (minx + maxx) / 2
        center_y = (miny + maxy) / 2
        half_tile = tile_size / 2
        win_minx = max(0, center_x - half_tile)
        win_miny = max(0, center_y - half_tile)
        win_maxx = min(image_width, center_x + half_tile)
        win_maxy = min(image_height, center_y + half_tile)

        # 调整以确保正好1024x1024（如果边界不足，padding或调整）
        if win_maxx - win_minx < tile_size:
            win_minx = max(0, win_maxx - tile_size)
        if win_maxy - win_miny < tile_size:
            win_miny = max(0, win_maxy - tile_size)

        proposed_window = Window(win_minx, win_miny, tile_size, tile_size)

        # 检查与现有窗口的重叠，确保<20%
        too_much_overlap = False
        for existing_win in generated_windows:
            # 计算重叠区域
            overlap_x = max(0, min(win_minx + tile_size, existing_win.col_off + tile_size) - max(win_minx, existing_win.col_off))
            overlap_y = max(0, min(win_miny + tile_size, existing_win.row_off + tile_size) - max(win_miny, existing_win.row_off))
            overlap_area = overlap_x * overlap_y
            if overlap_area > (tile_size * tile_size * max_overlap_ratio):
                too_much_overlap = True
                break

        if too_much_overlap:
            # 如果重叠太大，调整位置（这里简单跳过，您可实现更复杂的调整，如偏移步长）
            print(f"图斑 {idx} 的窗口重叠过多，跳过或调整")
            continue

        # 添加到列表
        generated_windows.append(proposed_window)

        # 裁剪并保存切片
        tile_data = src.read(window=proposed_window)
        tile_meta = src.meta.copy()
        tile_meta.update({
            'width': proposed_window.width,
            'height': proposed_window.height,
            'transform': rasterio.windows.transform(proposed_window, src.transform)
        })
        output_path = os.path.join(output_dir, f'tile_{idx}.tif')
        with rasterio.open(output_path, 'w', **tile_meta) as dst:
            dst.write(tile_data)

        print(f"生成切片 {output_path}，覆盖图斑 {idx}")

    print(f"总共生成 {len(generated_windows)} 个切片")

# 注意：这个代码确保每个切片包含至少一个完整图斑，但可能有图斑被多个切片覆盖（如果重叠允许）。如果需要覆盖所有图斑的最优 tiling，请提供更多细节。