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


def merge_and_rasterize_vectors_four_categories(file_paths, output_vector_path, output_raster_path,
                                                raster_resolution=30):
    """
    合并4个矢量图层（bareland, grass, shrub, building），转为栅格并显示4个类别，保留0值作为背景

    参数:
    file_paths: 包含4个矢量文件路径的列表 (按顺序: bareland, grass, shrub, building)
    output_vector_path: 输出合并后的矢量文件路径
    output_raster_path: 输出栅格文件路径
    raster_resolution: 输出栅格的分辨率（默认30米）
    """
    try:
        # 检查输入文件数量
        if len(file_paths) != 4:
            raise ValueError("请提供正好4个矢量文件路径")

        # 定义类别映射，按照指定顺序
        categories = ['bareland', 'grass', 'shrub', 'building']
        category_values = {'bareland': 1, 'grass': 2, 'shrub': 3, 'building': 4}  # bareland=1, grass=2, shrub=3, building=4

        # 读取所有矢量文件并添加分类
        gdfs = []
        for i, (path, category) in enumerate(zip(file_paths, categories), 1):
            print(f"正在读取第{i}个矢量文件: {path}")
            gdf = gpd.read_file(path)
            gdf['category'] = category  # 添加类别字段
            gdf['value'] = category_values[category]  # 添加数值字段用于栅格化
            gdfs.append(gdf)

        # 检查和统一坐标系（以第一个图层为基准）
        base_crs = gdfs[0].crs
        for i, gdf in enumerate(gdfs[1:], 2):
            if gdf.crs != base_crs:
                print(f"第{i}个图层坐标系不同，正在转换为第一个图层的坐标系...")
                gdfs[i - 1] = gdf.to_crs(base_crs)

        # 合并所有GeoDataFrame
        print("正在合并图层...")
        merged_gdf = gpd.GeoDataFrame(pd.concat(gdfs, ignore_index=True))
        merged_gdf = merged_gdf.set_geometry('geometry')

        # 保存合并后的矢量文件
        print(f"正在保存合并矢量结果到: {output_vector_path}")
        merged_gdf.to_file(output_vector_path)

        # 计算栅格化的范围
        bounds = merged_gdf.total_bounds  # [minx, miny, maxx, maxy]
        width = int((bounds[2] - bounds[0]) / raster_resolution)
        height = int((bounds[3] - bounds[1]) / raster_resolution)

        # 创建栅格化变换
        transform = rasterio.transform.from_bounds(
            bounds[0], bounds[1], bounds[2], bounds[3], width, height
        )

        # 栅格化：将几何对象转为栅格，使用'value'字段作为像素值
        print("正在将矢量转为栅格...")
        shapes = ((mapping(geom), value) for geom, value in zip(merged_gdf.geometry, merged_gdf['value']))
        raster = rasterize(
            shapes=shapes,
            out_shape=(height, width),
            transform=transform,
            fill=0,  # 背景值为0
            dtype=rasterio.uint8
        )

        # 保存栅格文件
        with rasterio.open(
                output_raster_path,
                'w',
                driver='GTiff',
                height=height,
                width=width,
                count=1,
                dtype=rasterio.uint8,
                crs=base_crs,
                transform=transform,
                nodata=0  # 设置nodata值为0，表示背景
        ) as dst:
            dst.write(raster, 1)

        print("处理完成！")
        print(f"合并矢量保存为: {output_vector_path}")
        print(f"栅格保存为: {output_raster_path}")
        print(f"栅格中类别值: 0=背景, 1=bareland, 2=grass, 3=shrub, 4=building")

    except Exception as e:
        print(f"发生错误: {str(e)}")


# 使用示例
if __name__ == "__main__":
    # 输入4个矢量文件路径（按顺序: bareland, grass, shrub, building）
    input_files = [
        r"E:\njds\shenyang\shengyangshp\bareland\bareland.shp",
        r"E:\njds\shenyang\shengyangshp\grass\grass2.shp",
        r"E:\njds\shenyang\shengyangshp\shrub\shrub.shp",
        r"E:\Data\z18\sd\Total_sd\shandonghebingshp\sdshp\jianzhu\true\shandongjianzhu3.shp"
    ]
    output_vector_file = r"E:\njds\shenyang\shengyangshp\kejianshe\merged_vector_four.shp"
    output_raster_file = r"E:\njds\shenyang\shengyangshp\kejianshe\merged_raster_four.tif"

    # 执行合并和栅格化
    merge_and_rasterize_vectors_four_categories(input_files, output_vector_file, output_raster_file,
                                                raster_resolution=10)