from AtmosphericDelayAuxData import NcHandle,ReadImage
from AtmosphericDelayAlg import AtmosphericDelay as Ady
from tool.algorithm.image import ImageHandle
from tool.algorithm.algtools.PreProcess import PreProcess as pp
import gc
import os
import sys
from pathlib import Path

env_str = os.path.dirname(os.path.abspath(sys.argv[0]))
os.environ['PROJ_LIB'] = env_str


def searchfile(source, pattern='*', recursive=False):

    if not isinstance(source, Path):
        source = Path(source)
    if source.is_dir() and source.exists():
        if recursive:
            res = list(source.rglob(pattern))
        else:
            res = list(source.glob(pattern))
    else:
        print(f"{source} is a file path or folder path is not exists!")
        return []   # False
    file_list = [i for i in res if i.is_file()]
    return file_list


def get_nc_dic(nc_dir):
    imageHandler = ImageHandle.ImageHandler()
    arr_dict = {}
    nc_list = ["AuxiliaryNC_geo_h", "AuxiliaryNC_re_hum", "AuxiliaryNC_temp",
               "MasterNC_geo_h", "MasterNC_re_hum", "MasterNC_temp"]
    arr_name = ["a_geo", "a_re_hum", "a_temp", "m_geo", "m_re_hum", "m_temp"]
    path = nc_dir
    for i in range(0, len(nc_list)):
        in_path = path + nc_list[i] + ".tif"
        array = imageHandler.get_all_band_array(in_path)
        arr_dict.update({arr_name[i]: array})
    return arr_dict["a_geo"], arr_dict["a_re_hum"], arr_dict["a_temp"], arr_dict["m_geo"], arr_dict["m_re_hum"], arr_dict["m_temp"]


def resampe_image(aux_dry_wet_unresame_path, mas_dry_wet_unresame_path, temp_dir):
    """
    将延迟图（ztd）恢复为非正方形的样子
    para: aux_dry_wet_unresame_path：待裁剪的辅影像延迟图路径
    para: mas_dry_wet_unresame_path：待裁剪的主影像延迟图路径
    return :裁剪后的辅影像延迟图路径、主影像延迟图路径
    """
    # 1.3 裁剪主辅影像延迟值
    out_aux_dry_wet_resamed_path = temp_dir + "resamed_aux_dry_wet.tif"
    out_mas_dry_wet_resamed_path = temp_dir + "resamed_mas_dry_wet.tif"

    pp.resampling_by_scale(aux_dry_wet_unresame_path, out_aux_dry_wet_resamed_path,
                           temp_dir + "MasterNC_geo_h.tif")
    pp.resampling_by_scale(mas_dry_wet_unresame_path, out_mas_dry_wet_resamed_path,
                           temp_dir + "MasterNC_geo_h.tif")

    return out_aux_dry_wet_resamed_path, out_mas_dry_wet_resamed_path


if __name__ == '__main__':
    mas_nc = r'D:\micro\WorkSpace\AtmosphericDelay\ERA5\ERA5_N42.0_N46.0_E114.0_E118.0_20230615_08.nc'
    aux_nc = r'D:\micro\WorkSpace\AtmosphericDelay\ERA5\ERA5_N42.0_N46.0_E114.0_E118.0_20230910_08.nc'
    dem_file = r'D:\micro\WorkSpace\AtmosphericDelay\ERA5\cut_dem.tif'
    temp_dir = r'D:\micro\WorkSpace\AtmosphericDelay\test_ztd\test'
    out_path = r'D:\micro\WorkSpace\AtmosphericDelay\test_ztd' + '\\'

    # 读取气象数据
    ncHandle = NcHandle()
    imageHandle = ImageHandle.ImageHandler()
    ncHandle.tran_tif(mas_nc, out_path, "MasterNC")
    ncHandle.tran_tif(aux_nc, out_path, "AuxiliaryNC")

    a_geo, a_re_hum, a_temp, m_geo, m_re_hum, m_temp = get_nc_dic(out_path)

    base_file = out_path + "MasterNC_geo_h.tif"  # 用来获取裁剪后气象数据的经纬度、分辨率等信息
    #
    # 1.2 计算主、辅影像垂直分层湿干延迟值，并保存干湿延迟图
    m_cdstack_dry_wet = Ady().aps_wrf_sar(base_file, m_temp, m_re_hum, m_geo, dem_file)  # 主影像垂直分层干湿延迟值数组
    m_dry_wet_path = out_path + "m_dry_wet.tif"  # 主影像干湿延迟图保存路径
    Ady().write_dry_wet_tif(dem_file, base_file, m_cdstack_dry_wet, m_dry_wet_path)

    a_cdstack_dry_wet = Ady().aps_wrf_sar(base_file, a_temp, a_re_hum, a_geo, dem_file)  # 辅影像垂直分层干湿延迟值数组
    a_dry_wet_path = out_path + "a_dry_wet.tif"  # 辅影像干湿延迟图保存路径
    Ady().write_dry_wet_tif(dem_file, base_file, a_cdstack_dry_wet, a_dry_wet_path)
    #
    # aux_dry_wet_resamed_path, mas_dry_wet_resamed_path = resampe_image(a_dry_wet_path, m_dry_wet_path, temp_dir)  # 参照气象数据分辨率进行重采样

    gc.collect()  # 回收内存
    # 1.4 输出主辅影像ztd数组
    a_ztd = Ady().calc_ztd(a_dry_wet_path, dem_file)  # 辅影像的ztd数组
    m_ztd = Ady().calc_ztd(m_dry_wet_path, dem_file)  # 主影像的ztd数组
    out_aux_ztd_path = out_path + "aux_ztd.tif"
    out_mas_ztd_path = out_path + "mas_ztd.tif"
    Ady().write_ztd_tif(dem_file, base_file, a_ztd, out_aux_ztd_path)  # ztd数组->ztd影像
    Ady().write_ztd_tif(dem_file, base_file, m_ztd, out_mas_ztd_path)  # ztd数组->ztd影像