lut.py

'''
Placeholder for model-based correction LUT
'''
import os
import isce3
import numpy as np
from osgeo import gdal
import pysolid
from scipy.interpolate import RegularGridInterpolator as RGI
from skimage.transform import resize

from compass.utils.geometry_utils import enu2los, en2az
from compass.utils.helpers import open_raster


def cumulative_correction_luts(burst, dem_path,
                               rg_step=200, az_step=0.25,
                               scratch_path=None):
    '''
    Sum correction LUTs and returns cumulative correction LUT in slant range
    and azimuth directions

    Parameters
    ----------
    burst: Sentinel1BurstSlc
        Sentinel-1 A/B burst SLC object
    dem_path: str
        Path to the DEM file
    rg_step: float
        LUT spacing along slant range direction
    az_step: float
        LUT spacing along azimuth direction
    scratch_path: str
        Path to the scratch directory

    Returns
    -------
    rg_lut: isce3.core.LUT2d
        Sum of slant range correction LUTs in meters as a function of azimuth
        time and slant range
    az_lut: isce3.core.LUT2d
        Sum of azimuth correction LUTs in seconds as a function of azimuth time
        and slant range
    '''
    # Get individual LUTs
    geometrical_steer_doppler, bistatic_delay, az_fm_mismatch, tides = \
        compute_geocoding_correction_luts(burst,
                                          dem_path=dem_path,
                                          rg_step=rg_step,
                                          az_step=az_step,
                                          scratch_path=scratch_path)

    # Convert to geometrical doppler from range time (seconds) to range (m)
    rg_lut_data = \
        geometrical_steer_doppler.data * isce3.core.speed_of_light * 0.5 + \
        tides[0]

    # Invert signs to correct for convention
    # TO DO: add azimuth SET to LUT
    az_lut_data = -(bistatic_delay.data + az_fm_mismatch.data)

    rg_lut = isce3.core.LUT2d(bistatic_delay.x_start,
                              bistatic_delay.y_start,
                              bistatic_delay.x_spacing,
                              bistatic_delay.y_spacing,
                              rg_lut_data)
    az_lut = isce3.core.LUT2d(bistatic_delay.x_start,
                              bistatic_delay.y_start,
                              bistatic_delay.x_spacing,
                              bistatic_delay.y_spacing,
                              az_lut_data)

    return rg_lut, az_lut


def compute_geocoding_correction_luts(burst, dem_path,
                                      rg_step=200, az_step=0.25,
                                      scratch_path=None):
    '''
    Compute slant range and azimuth LUTs corrections
    to be applied during burst geocoding
    Parameters
    ----------
    burst: Sentinel1BurstSlc
        S1-A/B burst object
    dem_path: str
        Path to the DEM required for azimuth FM rate mismatch.
    xstep: int
        LUT spacing along x/slant range in meters
    ystep: int
        LUT spacing along y/azimuth in seconds

    scratch_path: str
        Path to the scratch directory.
        If `None`, `burst.az_fm_rate_mismatch_mitigation()` will
        create temporary directory internally.

    Returns
    -------
    geometrical_steering_doppler: isce3.core.LUT2d:
        LUT2D object of total doppler (geometrical doppler +  steering doppler)
        in seconds as the function of the azimuth time and slant range.
        This correction needs to be added to the SLC tagged range time to
        get the corrected range times.

    bistatic_delay: isce3.core.LUT2d:
        LUT2D object of bistatic delay correction in seconds as a function
        of the azimuth time and slant range.
        This correction needs to be added to the SLC tagged azimuth time to
        get the corrected azimuth times.

    az_fm_mismatch: isce3.core.LUT2d:
        LUT2D object of azimuth FM rate mismatch mitigation,
        in seconds as the function of the azimuth time and slant range.
        This correction needs to be added to the SLC tagged azimuth time to
        get the corrected azimuth times.

    [rg_set, az_set]: list, np.ndarray
        List of numpy.ndarray containing SET in slant range and azimuth directions
        in meters. These corrections need to be added to the slC tagged azimuth
        and slant range times.
    '''
    # Get DEM raster
    dem_raster = isce3.io.Raster(dem_path)

    # Create directory to store SET temp results
    output_path = f'{scratch_path}/corrections'
    os.makedirs(output_path, exist_ok=True)

    # Compute Geometrical Steering Doppler
    geometrical_steering_doppler = \
        burst.doppler_induced_range_shift(range_step=rg_step, az_step=az_step)

    # Compute bistatic delay
    bistatic_delay = burst.bistatic_delay(range_step=rg_step, az_step=az_step)

    # Compute azimuth FM-rate mismatch
    if not os.path.isfile(dem_path):
        raise FileNotFoundError(f'Cannot find the dem file: {dem_path}')
    az_fm_mismatch = burst.az_fm_rate_mismatch_mitigation(dem_path,
                                                          scratch_path,
                                                          range_step=rg_step,
                                                          az_step=az_step)

    # Get solid Earth tides on a very coarse grid
    # Run rdr2geo on a very coarse grid
    lon_path, lat_path, hgt_path, inc_path, head_path = \
        compute_rdr2geo_rasters(burst, dem_raster, output_path,
                                rg_step * 10, az_step * 10)

    # Open rdr2geo layers and feed them to SET computation
    lat = open_raster(lat_path)
    lon = open_raster(lon_path)
    hgt = open_raster(hgt_path)
    inc_angle = open_raster(inc_path)
    head_angle = open_raster(head_path)

    # compute Solid Earth Tides (using pySolid)
    rg_set_temp, az_set_temp = solid_earth_tides(burst, lat, lon,
                                                 inc_angle, head_angle)
    

    epsg_dem = dem_raster.get_epsg()
    proj = isce3.core.make_projection(epsg_dem)
    ellipsoid = proj.ellipsoid
    rg_set_temp_2, az_set_temp_2 = solid_earth_tides_2(burst, lat, lon, hgt, ellipsoid)
    
    # Resize SET to the size of the correction grid
    out_shape = bistatic_delay.data.shape
    kwargs = dict(order=1, mode='edge', anti_aliasing=True,
                  preserve_range=True)
    rg_set = resize(rg_set_temp, out_shape, **kwargs)
    az_set = resize(az_set_temp, out_shape, **kwargs)

    return geometrical_steering_doppler, bistatic_delay, az_fm_mismatch, [
        rg_set, az_set]


def solid_earth_tides(burst, lat_radar_grid, lon_radar_grid, inc_angle,
                      head_angle):
    '''
    Compute displacement due to Solid Earth Tides (SET)
    in slant range and azimuth directions

    Parameters
    ---------
    burst: Sentinel1Slc
        S1-A/B burst object
    lat_radar_grid: np.ndarray
        Latitude array on burst radargrid
    lon_radar_grid: np.ndarray
        Longitude array on burst radargrid
    inc_angle: np.ndarray
        Incident angle raster in unit of degrees
    head_angle: np.ndaaray
        Heading angle raster in unit of degrees

    Returns
    ------
    rg_set: np.ndarray
        2D array with SET displacement along LOS in meters
    az_set: np.ndarray
        2D array with SET displacement along azimuth in meters
    '''

    # Extract top-left coordinates from burst polygon
    lon_min, lat_min, _, _ = burst.border[0].bounds

    # Generate the atr object to run pySolid. We compute SET on a
    # 2.5 km x 2.5 km coarse grid
    margin = 0.1
    lat_start = lat_min - margin
    lon_start = lon_min - margin

    atr = {
        'LENGTH': 25,
        'WIDTH': 100,
        'X_FIRST': lon_start,
        'Y_FIRST': lat_start,
        'X_STEP': 0.023,
        'Y_STEP': 0.023
    }

    # Run pySolid and get SET in ENU coordinate system
    (set_e,
     set_n,
     set_u) = pysolid.calc_solid_earth_tides_grid(burst.sensing_start, atr,
                                                  display=False, verbose=True)

    # Resample SET from geographical grid to radar grid
    # Generate the lat/lon arrays for the SET geogrid
    lat_geo_array = np.linspace(atr['Y_FIRST'],
                                lat_start + atr['Y_STEP'] * atr['LENGTH'],
                                num=atr['LENGTH'])
    lon_geo_array = np.linspace(atr['X_FIRST'],
                                lon_start + atr['X_STEP'] * atr['WIDTH'],
                                num=atr['WIDTH'])

    # Use scipy RGI to resample SET from geocoded to radar coordinates
    pts_src = (np.flipud(lat_geo_array), lon_geo_array)
    pts_dst = (lat_radar_grid.flatten(), lon_radar_grid.flatten())

    rdr_set_e = resample_set(set_e, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)
    rdr_set_n = resample_set(set_n, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)
    rdr_set_u = resample_set(set_u, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)

    # Convert SET from ENU to range/azimuth coordinates
    # Note: rdr2geo heading angle is measured wrt to the East and it is positive
    # anti-clockwise. To convert ENU to LOS, we need the azimuth angle which is
    # measured from the north and positive anti-clockwise
    # azimuth_angle = heading + 90
    set_rg = enu2los(rdr_set_e, rdr_set_n, rdr_set_u, inc_angle,
                     az_angle=head_angle + 90.0)
    set_az = en2az(rdr_set_e, rdr_set_n, head_angle - 90.0)

    return set_rg, set_az


def solid_earth_tides_2(burst, lat_radar_grid, lon_radar_grid, hgt_radar_grid, ellipsoid):
    '''
    Compute displacement due to Solid Earth Tides (SET)
    in slant range (meters) and azimuth directions (seconds)

    Parameters
    ---------
    burst: Sentinel1Slc
        S1-A/B burst object
    lat_radar_grid: np.ndarray
        Latitude array on burst radargrid
    lon_radar_grid: np.ndarray
        Longitude array on burst radargrid
    hgt_radar_grid: np.ndarray
        Height array on burst radargrid
    inc_angle: np.ndarray
        Incident angle raster in unit of degrees
    ellipsoid: isce3.core.Ellipsoid
        Ellipsoid definition from the projection of the source DEM

    Returns
    ------
    rg_set: np.ndarray
        2D array with SET displacement along LOS in meters
    az_set: np.ndarray
        2D array with SET displacement along azimuth in seconds
    '''

    # Extract top-left coordinates from burst polygon
    lon_min, lat_min, _, _ = burst.border[0].bounds

    # Generate the atr object to run pySolid. We compute SET on a
    # 2.5 km x 2.5 km coarse grid
    margin = 0.1
    lat_start = lat_min - margin
    lon_start = lon_min - margin

    atr = {
        'LENGTH': 25,
        'WIDTH': 100,
        'X_FIRST': lon_start,
        'Y_FIRST': lat_start,
        'X_STEP': 0.023,
        'Y_STEP': 0.023
    }

    # Run pySolid and get SET in ENU coordinate system
    (set_e,
     set_n,
     set_u) = pysolid.calc_solid_earth_tides_grid(burst.sensing_start, atr,
                                                  display=False, verbose=True)

    # Resample SET from geographical grid to radar grid
    # Generate the lat/lon arrays for the SET geogrid
    lat_geo_array = np.linspace(atr['Y_FIRST'],
                                lat_start + atr['Y_STEP'] * atr['LENGTH'],
                                num=atr['LENGTH'])
    lon_geo_array = np.linspace(atr['X_FIRST'],
                                lon_start + atr['X_STEP'] * atr['WIDTH'],
                                num=atr['WIDTH'])

    # Use scipy RGI to resample SET from geocoded to radar coordinates
    pts_src = (np.flipud(lat_geo_array), lon_geo_array)
    pts_dst = (lat_radar_grid.flatten(), lon_radar_grid.flatten())

    rdr_set_e = resample_set(set_e, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)
    rdr_set_n = resample_set(set_n, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)
    rdr_set_u = resample_set(set_u, pts_src, pts_dst).reshape(
        lat_radar_grid.shape)

    # Convert enu in radar grid into rg and az
    set_rg, set_az = enu_to_rgaz(burst.as_isce3_radargrid(),
                                 burst.orbit,
                                 isce3.core.LUT2d(),
                                 np.array([lon_radar_grid, lat_radar_grid, hgt_radar_grid]),
                                 np.array([rdr_set_e, rdr_set_n, rdr_set_u]),
                                 ellipsoid)

    return set_rg, set_az


def compute_rdr2geo_rasters(burst, dem_raster, output_path,
                            rg_step, az_step):
    '''
    Get latitude, longitude, incidence and
    azimuth angle on multi-looked radar grid

    Parameters
    ----------
    burst: Sentinel1Slc
        S1-A/B burst object
    dem_raster: isce3.io.Raster
        ISCE3 object including DEM raster
    output_path: str
        Path where to save output rasters
    rg_step: float
        Spacing of radar grid along slant range
    az_step: float
        Spacing of the radar grid along azimuth

    Returns
    -------
    x_path: str
        Path to longitude raster
    y_path: str
        Path to latitude raster
    inc_path: str
        Path to incidence angle raster
    head_path: str
        Path to heading angle raster
    '''

    # Some ancillary inputs
    epsg = dem_raster.get_epsg()
    proj = isce3.core.make_projection(epsg)
    ellipsoid = proj.ellipsoid

    # Get radar and doppler grid
    width_rdr_grid, length_rdr_grid = [vec.size for vec in
                                       burst._steps_to_vecs(rg_step, az_step)]

    rdr_grid = isce3.product.RadarGridParameters(
        burst.as_isce3_radargrid().sensing_start,
        burst.wavelength,
        1.0 / az_step,
        burst.starting_range,
        rg_step,
        isce3.core.LookSide.Right,
        length_rdr_grid,
        width_rdr_grid,
        burst.as_isce3_radargrid().ref_epoch
    )

    grid_doppler = isce3.core.LUT2d()

    # Initialize the rdr2geo object
    rdr2geo_obj = isce3.geometry.Rdr2Geo(rdr_grid, burst.orbit,
                                         ellipsoid, grid_doppler,
                                         threshold=1.0e8)

    # Get the rdr2geo raster needed for SET computation
    topo_output = {f'{output_path}/x.rdr': gdal.GDT_Float64,
                   f'{output_path}/y.rdr': gdal.GDT_Float64,
                   f'{output_path}/z.rdr': gdal.GDT_Float64,
                   f'{output_path}/incidence_angle.rdr': gdal.GDT_Float32,
                   f'{output_path}/heading_angle.rdr': gdal.GDT_Float32}
    raster_list = [
        isce3.io.Raster(fname, rdr_grid.width,
                        rdr_grid.length, 1, dtype, 'ENVI')
        for fname, dtype in topo_output.items()]
    x_raster, y_raster, z_raster, incidence_raster, heading_raster = raster_list

    # Run rdr2geo on coarse radar grid
    rdr2geo_obj.topo(dem_raster, x_raster, y_raster, z_raster,
                     incidence_angle_raster=incidence_raster,
                     heading_angle_raster=heading_raster)

    # Return file path to rdr2geo layers
    paths = list(topo_output.keys())
    return paths[0], paths[1], paths[2], paths[3], paths[4]


def resample_set(geo_tide, pts_src, pts_dest):
    '''
    Use scipy RegularGridInterpolator to resample geo_tide
    from a geographical to a radar grid

    Parameters
    ----------
    geo_tide: np.ndarray
        Tide displacement component on geographical grid
    pts_src: tuple of ndarray
        Points defining the source rectangular regular grid for resampling
    pts_dest: tuple of ndarray
        Points defining the destination grid for resampling
    Returns
    -------
    rdr_tide: np.ndarray
        Tide displacement component resampled on radar grid
    '''

    # Flip tide displacement component to be consistent with flipped latitudes
    geo_tide = np.flipud(geo_tide)
    rgi_func = RGI(pts_src, geo_tide, method='nearest',
                   bounds_error=False, fill_value=0)
    rdr_tide = rgi_func(pts_dest)
    return rdr_tide


def enu_to_rgaz(radargrid_ref, orbit, doppler, 
                arr_llh,
                arr_enu,
                ellipsoid):
    '''
    Convert ENU displacement into range / azimuth

    Parameters
    ----------
    radargrid_ref:
        Radargrid of the burst
    orbit: isce3.core.Orbit
        Orbit of the burst
    doppler: isce3.core.LUT2d
        Doppler to be used for geo2rdr
    arr_llh: numpy.ndarray
        3 * N * M array, with order of lon, lat, hgt
    arr_enu: numpy.ndarray
        3 * N * M array, with order of easting, northing, up
    ellipsoid: isce3.core.Ellipsoid
        Ellipsoid definition to retrieve the radius of the ellipsoid
    
    Returns
    -------
    arr_rg: np.ndarray
        N * M array as the displacements in range direction in meters
    arr_az: np.ndarray
        N * M array as the displacements in azimut direction in seconds

    '''
    
    arr_az = np.zeros(arr_llh.shape[1:])
    arr_rg = np.zeros(arr_llh.shape[1:])

    arr_llh_displaced = llh_displaced_enu(arr_llh, arr_enu, ellipsoid)

    for i in range(arr_llh.shape[1]):
        for j in range(arr_llh.shape[2]):
            lon_deg = arr_llh[0, i, j]
            lat_deg = arr_llh[1, i, j]
            hgt_m = arr_llh[2, i, j]
            llh_ref = np.array([np.deg2rad(lon_deg),
                                np.deg2rad(lat_deg),
                                hgt_m])

            aztime_ref, slant_range_ref =\
                isce3.geometry.geo2rdr(llh_ref,
                                       ellipsoid,
                                       orbit,
                                       doppler,
                                       radargrid_ref.wavelength,
                                       radargrid_ref.lookside,
                                       threshold=1.0e-10, maxiter=50, delta_range=10.0)

            lon_deg_displaced = arr_llh_displaced[0, i, j]
            lat_deg_displaced = arr_llh_displaced[1, i, j]
            hgt_m_displaced = arr_llh_displaced[2, i, j]
            llh_displaced = np.array([np.deg2rad(lon_deg_displaced),
                                      np.deg2rad(lat_deg_displaced),
                                      hgt_m_displaced])
            
            aztime_displaced, slant_range_displaced =\
                isce3.geometry.geo2rdr(llh_displaced,
                                       ellipsoid,
                                       orbit,
                                       doppler,
                                       radargrid_ref.wavelength,
                                       radargrid_ref.lookside,
                                       threshold=1.0e-10, maxiter=50, delta_range=10.0)

            arr_az[i, j] = aztime_displaced - aztime_ref
            arr_rg[i, j] = slant_range_displaced - slant_range_ref

    return arr_rg, arr_az


def llh_displaced_enu(llh_ref, enu, ellipsoid):
    '''
    returns lon / lat / hgt that east / north / up is added to `llh_ref`

    Parameters
    ----------
    llh: np.ndarray
        3 * N * M array, with order of lon, lat, hgt
    enu: numpy.ndarray
        3 * N * M array, with order of easting, northing, up
    ellipsoid: isce3.core.Ellipsoid
        Ellipsoid definition to retrieve the radius of the ellipsoid

    Returns:
    --------
    llh_displaced: np.ndarray
        3 * N * M array, with order of lon, lat, hgt,
        of the grids in `llf_ref` displaced by `enu`
    
    '''

    llh_displaced = np.zeros(llh_ref.shape)

    for i in range(llh_ref.shape[1]):
        for j in range(llh_ref.shape[2]):
            llh_ref_rad = np.array([np.deg2rad(llh_ref[0,i,j]),
                                    np.deg2rad(llh_ref[1,i,j]),
                                    llh_ref[2,i,j]])
            
                                   
            radian_north = enu[0, i, j] / ellipsoid.r_north(llh_ref_rad[0])
            radian_east = enu[1, i, j] / ellipsoid.r_east(llh_ref_rad[0]) # TODO discuss. I might be using the incorrect one

            llh_displaced[0, i, j] = llh_ref[0, i, j] + np.rad2deg(radian_north)
            llh_displaced[1, i, j] = llh_ref[1, i, j] + np.rad2deg(radian_east)
            llh_displaced[2, i, j] = llh_ref[2, i, j] + enu[2, i, j]
    
    return llh_displaced


def tropo_static_model(hgt, inc_angle):
    '''
    Calculate the troposperic delay using static model

    Parameters
    ----------
    hgt: np.ndarray
        height of the surface

    inc_angle: np.ndarray
        Incidence angle in degrees

    Returns
    -------
    tropo: Tropospheric delay in meters
    
    '''

    ZPD = 2.3
    H = 6000.0
    tropo = ZPD / np.cos(np.deg2rad(inc_angle)) * np.exp(-1 * hgt / H)

    return tropo