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

def nc_to_tiff(nc_file, output_dir, clip_extent):
    """
    将 NetCDF 数据按时间维度裁剪并逐日导出为 GeoTIFF。

    参数:
        nc_file (str): 输入 NetCDF 文件路径。
        output_dir (str): 输出 GeoTIFF 文件夹路径；若不存在会自动创建。
        clip_extent (tuple[float, float, float, float]): 裁剪经纬度范围，格式为
            (min_lon, max_lon, min_lat, max_lat)，单位为度，坐标系默认为 EPSG:4326。
    """
    os.makedirs(output_dir, exist_ok=True)

    min_lon, max_lon, min_lat, max_lat = clip_extent

    # 创建裁剪区域（矩形）
    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)
    progress_step = max(1, total_days // 10)

    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) % progress_step == 0 or (i + 1) == total_days:
            print(f"已完成 {i + 1}/{total_days} ({(i + 1) / total_days * 100:.1f}%)")

    print(f"转换完成！共生成 {total_days} 个 tiff 文件")
    print(f"输出目录: {output_dir}")


if __name__ == "__main__":
    nc_to_tiff(
        nc_file=r"CN05.1_Tmax_1961-2023_daily.nc",
        output_dir=r"tiff_output",
        clip_extent=(113, 120, 36, 43),
    )