import argparse
import os
from osgeo import gdal, osr
import numpy as np
os.environ['GDAL_CACHEMAX'] = '2048' 
def resample_to_match(image_A_path, image_B_path, output_path, resample_method=gdal.GRA_Bilinear):
    """
    将图像A重采样至与图像B相同的坐标网格和分辨率
    :param image_A_path: 待重采样的图像A路径
    :param image_B_path: 参考图像B路径
    :param output_path: 输出文件路径
    :param resample_method: 重采样方法（默认双线性插值）
    """
    gdal.AllRegister()
    # 1. 打开参考图像B，获取其空间参数
    ds_B = gdal.Open(image_B_path, gdal.GA_ReadOnly)
    if ds_B is None:
        raise RuntimeError("无法打开参考图像B")

    # 提取B的投影、仿射变换、尺寸和波段数
    proj_B = ds_B.GetProjection()
    geotrans_B = ds_B.GetGeoTransform()
    width_B = ds_B.RasterXSize
    height_B = ds_B.RasterYSize
    bands_B = ds_B.RasterCount
    data_type_B = ds_B.GetRasterBand(1).DataType


    ds_A = gdal.Open(image_A_path, gdal.GA_ReadOnly)
    if ds_A is None:
        raise RuntimeError("无法打开输入图像A")

    bands_A=ds_A.RasterCount
    # 2. 创建输出文件，完全复制B的空间参数
    driver = gdal.GetDriverByName('ENVI')
    ds_out = driver.Create(
        output_path,
        width_B, height_B,
        bands_A,
        data_type_B
     )
    ds_out.SetGeoTransform(geotrans_B)
    ds_out.SetProjection(proj_B)

    # 3. 执行重采样

    # 配置WarpOptions参数（关键修改点）
    warp_opts = gdal.WarpOptions(
        srcSRS=ds_A.GetProjection(),  # 输入投影
        dstSRS=proj_B,                # 目标投影
        resampleAlg=resample_method,  # 重采样方法（如gdal.GRA_Bilinear）
        warpMemoryLimit=2 * 1024**3,  # 单块内存限制为2GB（单位：字节）
        # 分块处理相关参数
        multithread=True,             # 启用多线程加速
    )
    
    # 执行Warp操作（自动分块处理）
    gdal.Warp(
        ds_out,      # 输出数据集
        ds_A,        # 输入数据集
        options=warp_opts
    )


    # 4. 释放资源
    ds_A = ds_B = ds_out = None
    print(f"重采样完成，结果已保存至: {output_path}")


    # 4. 释放资源
    ds_A = ds_B = ds_out = None
    print(f"重采样完成，结果已保存至: {output_path}")

def getParams():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--infile',default=r'F:\Data\SAR\IDL\SAR_EC_20240801\total_precipitation2024080110.bin', help='输入文件')
    parser.add_argument('-o', '--outfile',default=r'F:\Data\SAR\IDL\SAR_EC_20240801\total_precipitation2024080110_resample.bin', help='输出Bin文件路径')
    parser.add_argument('-r', '--reffile',default=r'F:\Data\SAR\IDL\SAR_EC_20240801\4254-DB-GEO.tif', help='参考影像')
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        '--nearest',
        action='store_const',
        const='nearest',
        dest='method',
        help='启用最邻近插值（默认）'
    )
    group.add_argument(
        '--bilinear',
        action='store_const',
        const='bilinear',
        dest='method',
        help='启用双线性插值'
    )
    parser.set_defaults(method='nearest')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    parser = getParams()
    intiffPath=parser.infile
    outbinPath=parser.outfile
    refbinPath=parser.reffile
    methodstr=parser.method
    print('infile=',intiffPath)
    print('outfile=',outbinPath)
    print('reffile=',refbinPath)
    print('method=',methodstr)

    resamplemethodgdal=None
    if methodstr == 'nearest':
        resamplemethodgdal=gdal.GRA_NearestNeighbour
    elif methodstr == 'bilinear':
        resamplemethodgdal=gdal.GRA_Bilinear
    else:
        resamplemethodgdal = gdal.GRA_NearestNeighbour

    resample_to_match(
        image_A_path=intiffPath,
        image_B_path=refbinPath,
        output_path=outbinPath,
        resample_method=resamplemethodgdal  # 可选：gdal.GRA_NearestNeighbour（最邻近）
    )
    print("resample to match successfully")