import numpy as np
import rasterio
from rasterio import features
import geopandas as gpd
from shapely.geometry import mapping

def shp_to_mask(shp_path, reference_image_path, output_mask_path, burn_value=1):
    """
    将 SHP 多边形文件转换为与参考图像对齐的掩膜 ENVI 文件（.dat + .hdr）。

    参数
    ----------
    shp_path : str
        输入 SHP 文件路径（多边形要素）
    reference_image_path : str
        参考高光谱图像路径（用于获取空间范围、分辨率、投影）
    output_mask_path : str
        输出掩膜文件路径（应包含 .dat 扩展名）
    burn_value : int, optional
        掩膜内部像素值（默认为1，外部为0）
    """
    # 1. 读取 SHP 文件
    gdf = gpd.read_file(shp_path)
    if not gdf.geometry.geom_type.isin(['Polygon', 'MultiPolygon']).all():
        raise ValueError("SHP 必须包含 Polygon 或 MultiPolygon 几何类型。")

    # 2. 打开参考图像获取元数据
    with rasterio.open(reference_image_path) as src:
        ref_crs = src.crs
        ref_transform = src.transform
        ref_width = src.width
        ref_height = src.height

        # 3. 坐标转换（如果需要）
        if ref_crs is not None and gdf.crs != ref_crs:
            gdf = gdf.to_crs(ref_crs)

        # 4. 生成掩膜数组
        geometries = [mapping(geom) for geom in gdf.geometry]
        mask = features.geometry_mask(
            geometries,
            out_shape=(ref_height, ref_width),
            transform=ref_transform,
            all_touched=False,      # 仅中心点在多边形内
            invert=True            # True=内部, False=外部（这里取反）
        )
        # 转换布尔为整数掩膜（True -> burn_value, False -> 0）
        mask = mask.astype(np.uint8) * burn_value

        # 5. 写入 ENVI 格式（.dat + .hdr）
        # 继承参考图像的元数据，修改为单波段 ENVI 格式
        profile = src.profile
        profile.update(
            driver='ENVI',          # 输出格式
            count=1,
            dtype='uint8',
            compress=None          # ENVI 不支持 LZW 压缩
        )
        # 确保输出文件扩展名为 .dat
        if not output_mask_path.lower().endswith('.dat'):
            output_mask_path += '.dat'
        with rasterio.open(output_mask_path, 'w', **profile) as dst:
            dst.write(mask, 1)

    print(f"掩膜已保存至: {output_mask_path} (及同名的 .hdr 头文件)")

# ==================== 使用示例 ====================
if __name__ == "__main__":
    shp_file = r"E:\is2\yaopu\925\roi.shp"
    reference_img = r"E:\is2\yaopu\output\before.dat"   # 参考图像（可以是任意一幅高光谱图像）
    output_mask = r"E:\is2\yaopu\925\roi_mask.dat"       # 输出的掩膜文件
    shp_to_mask(shp_file, reference_img, output_mask, burn_value=1)