import xarray as xr
import rioxarray
import geopandas as gpd
from shapely.geometry import box
import os
import pandas as pd

# 文件路径
nc_file = r"CN05.1_Tmax_1961-2023_daily.nc"
output_dir = r"tiff_output"

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

# 定义裁剪范围（经纬度）
min_lon, max_lon = 113, 120
min_lat, max_lat = 36, 43

# 创建裁剪区域（矩形）
bbox = box(min_lon, min_lat, max_lon, max_lat)
bbox_gdf = gpd.GeoDataFrame({'geometry': [bbox]}, crs='EPSG:4326')

# 打开数据集
ds = xr.open_dataset(nc_file)
# 为数据添加空间参考
ds = ds.rio.write_crs("EPSG:4326", inplace=True)

# 裁剪数据集
ds_clipped = ds.rio.clip(bbox_gdf.geometry, bbox_gdf.crs, drop=True)

# 按时间维度循环，每个时间步保存为一个tiff
total_days = len(ds_clipped.time.values)
for i, time in enumerate(ds_clipped.time.values):
    # 提取单个时间步的tmax数据
    da_time = ds_clipped['tmax'].sel(time=time)
    # 格式化时间字符串用于文件名
    time_str = pd.to_datetime(time).strftime('%Y%m%d')
    # 输出文件名
    output_file = os.path.join(output_dir, f"Tmax_{time_str}.tiff")
    # 保存为GeoTIFF
    da_time.rio.to_raster(output_file)
    # 每处理10%的文件打印一次进度
    if (i + 1) % (total_days // 10) == 0:
        print(f"已完成 {i + 1}/{total_days} ({(i + 1) / total_days * 100:.1f}%)")
print(f"转换完成！共生成 {total_days} 个tiff文件")
print(f"输出目录: {output_dir}")