# -*- coding: utf-8 -*-

"""
@File    : 17_栅格转面&面转栅格.py
@Author  : fungis@163.com
@Time    : 2023/06/08 19:41
@notice  :
"""

from osgeo import ogr, osr, gdal


def shape_to_raster(shapefile, rasterfile, savefile):
    """
    面转栅格TIF
    :param shapefile: 文件矢量Shapefile文件
    :param rasterfile: 栅格参照文件
    :param savefile: 文件输出路径
    :return:
    """
    data = gdal.Open(rasterfile, gdal.GA_ReadOnly)
    x_res = data.RasterXSize
    y_res = data.RasterYSize
    shape = ogr.Open(shapefile)
    layer = shape.GetLayer()
    targetDataset = gdal.GetDriverByName('GTiff').Create(savefile, x_res, y_res, 1, gdal.GDT_Float32)
    targetDataset.SetGeoTransform(data.GetGeoTransform())
    targetDataset.SetProjection(data.GetProjection())
    band = targetDataset.GetRasterBand(1)
    band.SetNoDataValue(-9999)
    band.FlushCache()
    gdal.RasterizeLayer(targetDataset, [1, 2, 3], layer)


def raster_to_shape(rasterfile, shapefile):
    """
    栅格TIF转面
    :param rasterfile: 栅格TIF输入文件
    :param shapefile: 文件输出路径-矢量Shapefile文件
    :return:
    """
    data = gdal.Open(rasterfile, gdal.GA_ReadOnly)
    inband = data.GetRasterBand(1)
    drv = ogr.GetDriverByName('ESRI Shapefile')
    polygon = drv.CreateDataSource(shapefile)
    prj = osr.SpatialReference()
    prj.ImportFromWkt(data.GetProjection())  ## 使用栅格的投影信息
    polygon_layer = polygon.CreateLayer(shapefile, srs=prj, geom_type=ogr.wkbMultiPolygon)
    newField = ogr.FieldDefn('Value', ogr.OFTInteger)
    polygon_layer.CreateField(newField)
    gdal.FPolygonize(inband, None, polygon_layer, 0)


if __name__ == '__main__':
    # 输入栅格(不支持map文件)
    input_raster = r'G:\pythonCode\env\python-gdal-test\data-use\tif\global\coast_distance.tif'
    # 输入要素类
    input_shp = r'G:\pythonCode\env\python-gdal-test\demo-time-axis\data-use\shp\Vic_polygon\Basin_boundary.shp'
    # 栅格转面
    raster_to_shape(input_raster, './results/test.shp')