"""
@Project ：windProject
@File ：WindSpeedInversion.py
@Function ：单景处理影像
@Author ：CZH
@Date ：2025/10/19 14:39
@Version ：1.0.0
"""

import os
import numpy as np

from ImageHandle import ImageHandler
from osgeo import ogr, gdal
import os
import argparse
import numpy as np
from PIL import Image
import math
from pathlib import Path
from WindSpeedModel import python_windspeedCmod5n,python_ERA5_windAngle

def read_img(filename):
    """
    影像读取
    :param filename:
    :return:
    """
    gdal.AllRegister()
    img_dataset = gdal.Open(filename)  # 打开文件

    if img_dataset is None:
        msg = 'Could not open ' + filename
        print("ERROR", msg)
        return None, None, None

    im_proj = img_dataset.GetProjection()  # 地图投影信息
    if im_proj is None:
        return None, None, None
    im_geotrans = img_dataset.GetGeoTransform()  # 仿射矩阵

    im_width = img_dataset.RasterXSize  # 栅格矩阵的行数
    im_height = img_dataset.RasterYSize  # 栅格矩阵的行数
    im_arr = img_dataset.ReadAsArray(0, 0, im_width, im_height)
    del img_dataset
    return im_proj, im_geotrans, im_arr


def write_img_envi(filename, im_proj, im_geotrans, im_data, no_data='null'):
    """
    影像保存
    :param filename: 保存的路径
    :param im_proj:
    :param im_geotrans:
    :param im_data:
    :param no_data: 把无效值设置为 nodata
    :return:
    """
    gdal_dtypes = {
        'int8': gdal.GDT_Byte,
        'unit16': gdal.GDT_UInt16,
        'int16': gdal.GDT_Int16,
        'unit32': gdal.GDT_UInt32,
        'int32': gdal.GDT_Int32,
        'float32': gdal.GDT_Float32,
        'float64': gdal.GDT_Float64,
    }
    if not gdal_dtypes.get(im_data.dtype.name, None) is None:
        datatype = gdal_dtypes[im_data.dtype.name]
    else:
        datatype = gdal.GDT_Float32
    # 判读数组维数
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    else:
        im_bands, (im_height, im_width) = 1, im_data.shape
    # 创建文件
    if os.path.exists(os.path.split(filename)[0]) is False:
        os.makedirs(os.path.split(filename)[0])
    driver = gdal.GetDriverByName("ENVI")  # 数据类型必须有，因为要计算需要多大内存空间
    dataset = driver.Create(filename, im_width, im_height, im_bands, datatype)
    dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
    dataset.SetProjection(im_proj)  # 写入投影
    if im_bands == 1:
        # outRaster.GetRasterBand(1).WriteArray(array)  # 写入数组数据
        outband = dataset.GetRasterBand(1)
        outband.WriteArray(im_data)
        if no_data != 'null':
            outband.SetNoDataValue(no_data)
        outband.FlushCache()
    else:
        for i in range(im_bands):
            outband = dataset.GetRasterBand(1 + i)
            outband.WriteArray(im_data[i])
            outband.FlushCache()
            # outRaster.GetRasterBand(i + 1).WriteArray(array[i])
    del dataset
    dataset=None

def windVector2windDir(in_wind_dir_path):
    winddir_proj, winddir_geotrans, winddir_arr = read_img(in_wind_dir_path)
    wind_angle=python_ERA5_windAngle(winddir_arr[0,:,:],winddir_arr[1,:,:])
    return wind_angle


def windSpeedInversionProcess(in_wind_dir_path,in_inc_angle_path,in_sar_sigma0_path,out_wind_speed_path):

    wind_angle = windVector2windDir(in_wind_dir_path)
    incangle_proj, incangle_geotrans, incangle_arr=read_img(in_inc_angle_path)
    sarsigma0_proj, sarsigma0_geotrans, sarsigma0_arr=read_img(in_sar_sigma0_path)

    windspeed_arr=python_windspeedCmod5n(wind_angle,incangle_arr,sarsigma0_arr)
    write_img_envi(out_wind_speed_path, sarsigma0_proj, sarsigma0_geotrans, windspeed_arr)

def parse_args():
    """
    解析参数文件
    :return:
    """
    parser = argparse.ArgumentParser(description='wind speed inversion')
    parser.add_argument('--input_wind_dir', type=str, default=r'D:\Programs\SpacetyLabelAIPlante\Manual-Labeling-Tool\WindSpeedModel\TestData\SAR_EC_20240801\total_precipitation2024080110_resample.bin', help='wind direct')
    parser.add_argument('--input_inc_angle', type=str, default=r'D:\Programs\SpacetyLabelAIPlante\Manual-Labeling-Tool\WindSpeedModel\TestData\SAR_EC_20240801\4254-DB-IncidenceAngle-GEO.bin', help='input incidence angle')
    parser.add_argument('--input_sar_sigma0', type=str, default=r'D:\Programs\SpacetyLabelAIPlante\Manual-Labeling-Tool\WindSpeedModel\TestData\SAR_EC_20240801\4254-DB-GEO.bin', help='input sar sigma0 (dB)')
    parser.add_argument('--output_wind_speed', type=str, default=r'D:\Programs\SpacetyLabelAIPlante\Manual-Labeling-Tool\WindSpeedModel\TestData\InverserWindSpeed.bin', help='output wind speed')

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args=parse_args()
    in_wind_dir_path = args.input_wind_dir
    in_inc_angle_path=args.input_inc_angle
    in_sar_sigma0_path = args.input_sar_sigma0
    out_wind_speed_path = args.output_wind_speed
    windSpeedInversionProcess(in_wind_dir_path,in_inc_angle_path,in_sar_sigma0_path,out_wind_speed_path)