"""
Author: Wanxue Zhu
Time: 2025 - April - 18
Description: 这个代码用于将csv的文件转化为geotif;
             csv的文件格式为，第一列为long经度；
             第二列为lat纬度，第三列为数值。

"""
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
import rasterio
from rasterio.features import rasterize
from rasterio.transform import from_origin
import numpy as np
import os


#------------------1. 将csv转为点shape矢量数据，保存----------------------------

def csv_to_shapefile(csv_file, shapefile_output, directory):
    """
    将CSV文件转换为Shapefile文件，并保存到指定目录。
    csv_file (str): CSV文件的路径。
    shapefile_output (str): 输出Shapefile文件的路径。
    directory (str): 设置工作目录路径。
    """
    df = pd.read_csv(csv_file)
    lon = df.iloc[:, 0]
    lat = df.iloc[:, 1]
    val = df.iloc[:, 2]
    geometry = [Point(xy) for xy in zip(lon, lat)]
    gdf = gpd.GeoDataFrame(df, geometry=geometry, crs="EPSG:4326")
    gdf.to_file(shapefile_output, driver="ESRI Shapefile")


# # 调用函数
# csv_file = r'D:\17417\Documents\backend\Water\Data\matched_data.csv'  # 请替换为你的CSV文件路径
# shapefile_output = r'..\Raster\水浇地.shp'  # 输出Shapefile路径
# directory = r"D:\17417\Documents\backend\Water\Data"  # 工作目录路径
#
# csv_to_shapefile(csv_file, shapefile_output, directory)


# -----------------2. 将点shapefile转化为Geotif ----------------------------------

def vector_to_raster(input_shapefile, template_tif, output_tif, field):
    """
    将点矢量数据转换为栅格数据，输出与模板tif文件一致的GeoTIFF文件。
    input_shapefile (str): 输入点矢量数据的Shapefile文件路径。
    template_tif (str): 用作模板的GeoTIFF文件路径。
    output_tif (str): 输出栅格化后的GeoTIFF文件路径。
    field (str): 用于栅格化的属性字段名。
    """
    gdf = gpd.read_file(input_shapefile)
    with rasterio.open(template_tif) as src:# 打开模板栅格，获取元信息
        template_meta = src.meta.copy()
        transform = src.transform
        width = src.width
        height = src.height
        crs = src.crs
    # 将矢量数据投影到模板的坐标系
    if gdf.crs != crs:
        gdf = gdf.to_crs(crs)
    # 栅格化
    shapes = ((geom, value) for geom, value in zip(gdf.geometry, gdf[field]))
    raster = rasterize(
        shapes=shapes,
        out_shape=(height, width),
        transform=transform,
        fill=np.nan,
        dtype='float32'
    )
    # 创建输出目录（若不存在）
    os.makedirs(os.path.dirname(output_tif), exist_ok=True)
    # 写入 GeoTIFF，与模板完全一致
    template_meta.update({
        "count": 1,
        "dtype": 'float32',
        "nodata": np.nan
    })
    with rasterio.open(output_tif, 'w', **template_meta) as dst:
        dst.write(raster, 1)
    print("Save successfully：", output_tif)

# # 调用函数
# input_shapefile = r"..\Raster\水浇地.shp"
# template_tif = r"..\Raster\meanTemp.tif"  # 用作模板的 tif
# output_tif = r"..\Raster\水浇地.tif"
# field = 'Prediction'
#
# vector_to_raster(input_shapefile, template_tif, output_tif, field)