import os
import multiprocessing
import numpy as np
import pandas as pd
import xarray as xr
from datetime import datetime
from dateutil.relativedelta import relativedelta
from glob import glob
from scipy.ndimage import binary_dilation, distance_transform_edt
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    from osgeo import gdal, ogr, osr
except ImportError:
    import gdal
    import ogr
    import osr

# 环境变量
os.environ['PROJ_LIB'] = r'D:\ProgramData\Miniconda3_py39\envs\geovis\Lib\site-packages\osgeo\data\proj'
os.environ['GDAL_DATA'] = r'D:\ProgramData\Miniconda3_py39\envs\geovis\Lib\site-packages\osgeo\data\gdal'


def exportTif(filePath, lonData, latData, minLon, maxLat, res, bandDataList, epsgCode=4326, valueType=gdal.GDT_Float32,
              noData=-999):
    '''
    输出tif文件
    '''
    try:
        # 判断目录是否存在
        fileFolder = os.path.dirname(filePath)
        if not os.path.isdir(fileFolder):
            os.makedirs(fileFolder)

        # 目标坐标系
        dstSRS = osr.SpatialReference()
        dstSRS.ImportFromEPSG(epsgCode)
        spatialRef = dstSRS.ExportToWkt()

        # 创建geotransform
        geotransform = (minLon, res, 0, maxLat, 0, -res)

        # 获取数据维度
        bandNnum, rows, cols = bandDataList.shape

        # 创建输出文件
        driver = gdal.GetDriverByName('GTiff')
        outdataset = driver.Create(filePath, cols, rows, bandNnum, valueType,
                                   options=['TILED=YES', 'COMPRESS=LZW', 'BIGTIFF=YES'])
        outdataset.SetGeoTransform(geotransform)
        outdataset.SetProjection(spatialRef)

        # 写入每个波段
        for bandIndex in range(bandNnum):
            outband = outdataset.GetRasterBand(bandIndex + 1)
            outband.WriteArray(bandDataList[bandIndex])
            outband.SetNoDataValue(noData)
            outband.ComputeStatistics(False)

        # 创建金字塔
        outdataset.BuildOverviews('average', [2, 4, 8, 16])

        # 清理资源
        del outband, outdataset
        return 1
    except Exception as e:
        print(f"导出TIF文件时出错: {str(e)}")
        return 0


def optimize_fill_process(varData, timeData, latData, lonData, noData=-9999, expand_pixels=20):
    """
    优化后的填充处理函数：
    1. 减少滚动计算迭代次数和窗口大小
    2. 使用静态陆地掩膜（假设不随时间变化）
    3. 并行处理时间切片
    """
    try:
        # 1. 创建DataArray
        var_dataarray = xr.DataArray(
            data=varData,
            dims=['time', 'latitude', 'longitude'],
            coords={
                'time': timeData,
                'latitude': latData,
                'longitude': lonData,
            }
        )

        # 2. 使用chunk方法创建分块数组（优化分块策略）
        chunks = {
            'time': 8,  # 增加时间分块大小
            'latitude': 170,
            'longitude': 360
        }
        var_dataarray = var_dataarray.chunk(chunks)

        # 3. 初步填充（局部插值）
        var_filled = var_dataarray.where(var_dataarray != noData)

        # 4. 优化滚动计算 - 小范围滚动（3x3窗口）
        print("开始小范围滚动计算...")
        window_size_small = 3
        neighbor_mean = var_filled.rolling(
            longitude=window_size_small,
            latitude=window_size_small,
            center=True,
            min_periods=1
        ).mean()
        var_filled = var_filled.fillna(neighbor_mean)

        # 5. 优化滚动计算 - 中等范围滚动（11x11窗口）
        print("开始中等范围滚动计算...")
        window_size_medium = 7  # 减小窗口大小
        neighbor_mean = var_filled.rolling(
            longitude=window_size_medium,
            latitude=window_size_medium,
            center=True,
            min_periods=1
        ).mean()
        var_filled = var_filled.fillna(neighbor_mean)

        # 6. 计算前面所有延迟计算的任务
        print("开始计算前连续插值结果...")
        filled_data = var_filled.compute()  # shape: [time, lat, lon]

        # 7. 使用静态陆地掩膜（假设不随时间变化）
        print("创建静态陆地掩膜...")
        # 使用第一个时间切片创建静态掩膜
        static_mask = (filled_data[0] == noData)
        # 膨胀操作创建扩展区域
        expanded_static_mask = binary_dilation(static_mask, iterations=expand_pixels)
        # 计算静态掩膜的欧氏距离变换
        _, static_indices = distance_transform_edt(static_mask, return_indices=True)

        # 8. 并行处理每个时间切片
        print("开始并行边界扩充操作...")
        ntime = filled_data.shape[0]
        filled_data_extended = np.empty_like(filled_data)

        # 使用线程池并行处理
        with ThreadPoolExecutor(max_workers=min(16, multiprocessing.cpu_count())) as executor:
            futures = {}
            for t in range(ntime):
                future = executor.submit(
                    process_time_slice,
                    filled_data[t],
                    static_mask,
                    expanded_static_mask,
                    static_indices,
                    noData
                )
                futures[future] = t

            for future in as_completed(futures):
                t = futures[future]
                filled_data_extended[t] = future.result()

        print("边界扩充操作完成。")
        return filled_data_extended

    except Exception as e:
        print(f"优化处理过程中出错: {str(e)}")
        raise


def process_time_slice(data_slice, static_mask, expanded_static_mask, static_indices, noData):
    """处理单个时间切片的函数，用于并行执行"""
    try:
        # 复制数据以避免修改原始数据
        slice_data = data_slice.copy()

        # 确定需要填充的区域（扩展的静态掩膜区域中当前为无效值的部分）
        fill_mask = expanded_static_mask & (slice_data == noData)

        # 如果存在需要填充的区域
        if np.any(fill_mask):
            # 使用预计算的静态索引进行填充
            slice_data[fill_mask] = slice_data[
                static_indices[0][fill_mask],
                static_indices[1][fill_mask]
            ]

        return slice_data
    except Exception as e:
        print(f"处理时间切片时出错: {str(e)}")
        return data_slice  # 出错时返回原始数据


def process_variable(ds, varName, elementName, timeData, lonData, latData, minLon, maxLat, dataRes,
                     productFolder, process_thread, epsgCode, valueType, noData):
    """
    处理单个变量的函数
    """
    try:
        print(f"开始处理变量: {varName}")

        # 获取变量数据
        varData = ds[varName].data

        # 优化填充处理，包括局部插值和边界扩充
        print(f"开始优化填充处理: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        optimized_varData = optimize_fill_process(varData, timeData, latData, lonData, noData)
        print(f"完成优化填充处理: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

        # 获取处理后的数据
        varData = optimized_varData  # shape: [time, lat, lon]

        # 多线程处理导出TIFF
        if process_thread > 1:
            pool = multiprocessing.Pool(processes=process_thread)

        # 循环时间维度
        for i in range(len(timeData)):
            foreTime = pd.to_datetime(timeData[i]).strftime('%Y%m%d%H')
            elementData = varData[i]

            # 输出文件路径
            fileName = f'{elementName}-{foreTime}.tif'
            filePath = os.path.join(productFolder, fileName)

            # 多线程或单线程处理
            if process_thread > 1:
                pool.apply_async(exportTif,
                                 (filePath, lonData, latData, minLon, maxLat, dataRes,
                                  np.array([elementData]), epsgCode, valueType, noData))
            else:
                ret = exportTif(filePath, lonData, latData, minLon, maxLat, dataRes,
                                np.array([elementData]), epsgCode, valueType, noData)

        # 关闭多线程池
        if process_thread > 1:
            pool.close()
            pool.join()

        print(f"完成变量处理: {varName}")

    except Exception as e:
        print(f"处理变量 {varName} 时出错: {str(e)}")


def main():
    try:
        # 开始时间
        startTime = datetime.now()
        print(f"程序开始运行: {startTime.strftime('%Y-%m-%d %H:%M:%S')}")

        # 参数设置
        category = 'wave'
        fileFolder = r'C:\Users\mac\Desktop\fill_alg_test_data\input'
        process_thread = 8
        epsgCode = 4326
        valueType = gdal.GDT_Float32
        noData = -9999

        # 空间参数
        minLon, maxLat = -180, 90
        dataRes = 1 / 12

        # 获取文件列表
        fileList = sorted(glob(os.path.join(fileFolder, "*.nc")), reverse=False)
        print(f"找到 {len(fileList)} 个输入文件")

        # 1. 读取文件
        print("开始读取文件...")
        ds = xr.open_mfdataset(fileList)

        # 排序
        ds = ds.sortby(ds.time)
        ds = ds.sortby(ds.latitude, ascending=False)
        ds = ds.sortby(ds.longitude)

        # 经纬度数据
        lonData = ds.longitude.data
        latData = ds.latitude.data

        # 2. 时间重采样（不同品类可能有差别）
        print("开始时间重采样...")
        if category == 'atmos':
            date_start = pd.to_datetime(ds['time'].data[0]).strftime('%Y%m%d%H%M%S')
            date_end = (pd.to_datetime(ds['time'].data[-1]) + relativedelta(hours=23)).strftime('%Y%m%d%H%M%S')
            timeData = pd.date_range(date_start, date_end, freq='H')
        elif category == 'ocean':
            date_start = (pd.to_datetime(ds['time'].data[0]) - relativedelta(hours=0.5)).strftime('%Y%m%d%H%M%S')
            date_end = (pd.to_datetime(ds['time'].data[-1]) - relativedelta(hours=0.5)).strftime('%Y%m%d%H%M%S')
            timeData = pd.date_range(date_start, date_end, freq='H')
            ds = ds.squeeze('depth')
        elif category == 'wave':
            date_start = (pd.to_datetime(ds['time'].data[0]) - relativedelta(hours=2)).strftime('%Y%m%d%H%M%S')
            date_end = pd.to_datetime(ds['time'].data[-1]).strftime('%Y%m%d%H%M%S')
            timeData = pd.date_range(date_start, date_end, freq='h')

        # 如果需要时间插值
        if category in ['ocean', 'wave']:
            print("开始时间插值...")
            ds = ds.interp(time=timeData, kwargs={'fill_value': 'extrapolate'})

        # 产品目录
        productFolder = r'C:\Users\mac\Desktop\fill_alg_test_data\height_output_optimized_2'
        if not os.path.isdir(productFolder):
            os.makedirs(productFolder)

        # 要处理的变量列表
        varNameList = ['VHM0']
        elementNameList = ['wave_height', 'wave_dir', 'wave_period']

        # 处理每个变量
        for varName, elementName in zip(varNameList, elementNameList):
            if varName in ds.data_vars:
                process_variable(
                    ds, varName, elementName, timeData, lonData, latData,
                    minLon, maxLat, dataRes, productFolder, process_thread,
                    epsgCode, valueType, noData
                )
            else:
                print(f"数据文件中不包含元素: {varName}")

        # 结束时间
        endTime = datetime.now()
        print(f"程序运行完成: {endTime.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"总耗时: {(endTime - startTime).seconds} 秒")

    except Exception as e:
        print(f"主程序出错: {str(e)}")
        raise


if __name__ == '__main__':
    main()

【这是我目前的代码，算法执行需要515s，我的终极目的就是实现边界填充20个像元，有没有什么办法能够加速，可以损失一些精度】