# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/CLI/load.ipynb.

# %% auto 0
__all__ = ['read_gamma_image', 'de_load_gamma_flatten_rslc', 'de_load_gamma_lat_lon', 'de_load_gamma_incidence_angle',
           'de_load_gamma_hgt', 'de_load_gamma_range', 'de_load_gamma_metadata']

# %% ../../nbs/CLI/load.ipynb 3
import glob
from pathlib import Path
import tempfile
import re
import os

import toml
import zarr
import numpy as np
import pandas as pd
from scipy.constants import speed_of_light
import dask
from dask import array as da
from dask import delayed
from dask.distributed import Client, LocalCluster

from .utils.logging import get_logger, log_args

from fastcore.script import call_parse

# %% ../../nbs/CLI/load.ipynb 5
def _rdc_width_nlines(image_par):
    """get slc width and number of lines.
    """
    with open(image_par) as f:
        for line in f:
            if re.search('range_samples',line):
                rdc_width = int(line.split()[1])
            if re.search('azimuth_lines',line):
                rdc_nlines = int(line.split()[1])
    return rdc_width, rdc_nlines

# %% ../../nbs/CLI/load.ipynb 6
def _geo_width_nlines(dem_par):
    """get dem width and number of lines.
    """
    with open(dem_par) as f:
        for line in f:
            if re.search('width',line):
                geo_width = int(line.split()[1])
            if re.search('nlines',line):
                geo_nlines = int(line.split()[1])
    return geo_width, geo_nlines

# %% ../../nbs/CLI/load.ipynb 7
def _cor_pos_dem(dem_par):
    """get corner lat and lon and post lat and lon
    """
    with open(dem_par) as f:
        for line in f:
            if re.search('corner_lat',line):
                cor_lat = float(line.split()[1])
            if re.search('corner_lon',line):
                cor_lon = float(line.split()[1])
            if re.search('post_lat',line):
                pos_lat = float(line.split()[1])
            if re.search('post_lon',line):
                pos_lon = float(line.split()[1])
    return cor_lat, cor_lon, pos_lat, pos_lon

# %% ../../nbs/CLI/load.ipynb 8
def _fetch_slc_par_date(rslc_dir,# str / Path
                       ):
    rslc_dir = Path(rslc_dir)
    rslcs = []
    rslc_pars = []
    dates = []
    for rslc in sorted(rslc_dir.glob('*.rslc')):
        rslc_par = rslc.parent / (rslc.name + '.par')
        assert rslc_par.exists(), f'{str(rslc_par)} not exists!'
        date = rslc.stem

        rslcs.append(rslc)
        rslc_pars.append(rslc_par)
        dates.append(date)
    rslcs_df = pd.DataFrame({'date':dates,'rslc':rslcs,'par':rslc_pars})
    return rslcs_df

# %% ../../nbs/CLI/load.ipynb 9
def read_gamma_image(imag:str, # gamma raster data
                     width:int, # data width
                     dtype:str='float', # data format, only 'float' and 'fcomplex' are supported
                    ):
    '''read gamma image into numpy array.'''
    if dtype == 'float':
        dt = '>f4'
    elif dtype == 'fcomplex':
        dt = '>c8'
    else: raise ValueError('Unsupported data type')

    datf = open(imag,"r")
    datf.seek(0,os.SEEK_END)
    size = datf.tell()
    nlines = int(size / int(dt[-1]) / width)
    data = np.fromfile(imag,dtype=dt,count=width*nlines)
    return data.reshape(-1,width)

# %% ../../nbs/CLI/load.ipynb 12
@call_parse
@log_args
def de_load_gamma_flatten_rslc(rslc_dir:str, # gamma rslc directory, the name of the rslc and their par files should be '????????.rslc' and '????????.rslc.par'
                               reference:str, # reference date, eg: '20200202'
                               hgt:str, # the DEM in radar coordinate
                               rslc_zarr:str, # output, the flattened rslcs stack in zarr format
                               az_chunk_size:int=-1, # rslcs stack azimuth chunk size, azimuth number of lines by default (one chunk)
                               log:str=None, # logfile, no log by default
                              ):
    '''Generate flatten rslc data from gamma command and convert them into zarr format.
    The shape of hgt should be same as one rslc image, i.e. the hgt file is generated with 1 by 1 look geocoding.
    '''
    logger = get_logger(logfile=log)
    rslcs = _fetch_slc_par_date(rslc_dir)
    with pd.option_context('display.max_colwidth', 0):
        logger.info('rslc found: \n'+str(rslcs))
    rslc_pars = rslcs['par'].to_list()
    dates = rslcs['date'].to_list()
    rslcs = rslcs['rslc'].to_list()

    reference_idx = dates.index(reference)
    ref_rslc = rslcs[reference_idx]
    ref_rslc_par = rslc_pars[reference_idx]
    hgt = Path(hgt).absolute()

    n_image = len(rslcs)
    width,nlines = _rdc_width_nlines(ref_rslc_par)
    logger.info(f'number of images: {n_image}.')
    logger.info(f'image number of lines: {nlines}.')
    logger.info(f'image width: {width}.')

    logger.info('run gamma command to generate flatten rslcs:')
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        sim_orbs = []
        for i,(date,rslc,rslc_par) in enumerate(zip(dates,rslcs,rslc_pars)):
            off_par = temp_dir/(reference+'_'+date+'.off')
            create_offset_command = f'create_offset {str(ref_rslc_par)} {str(rslc_par)} {str(off_par)} 1 1 1 0'
            logger.info('run command: ' + create_offset_command)
            os.system(create_offset_command)
            # pg.create_offset(ref_rslc_par,rslc_par,off_par,1,1,1,0)
            sim_orb = temp_dir/(reference+'_'+date+'.sim_orb')
            phase_sim_orb_command = f'phase_sim_orb {str(ref_rslc_par)} {str(rslc_par)} {str(off_par)} {str(hgt)} {str(sim_orb)} {str(ref_rslc_par)} - - 1 1'
            logger.info('run command: ' + phase_sim_orb_command)
            os.system(phase_sim_orb_command)
            # pg.phase_sim_orb(ref_rslc_par,rslc_par,off_par,hgt,sim_orb,ref_rslc_par,'-','-',1,1)
            sim_orbs.append(sim_orb)
        logger.info('flatten rslcs generated.')
        logger.info('using dask to load the flatten rslcs from gamma binary format to zarr.')
        logger.info('starting dask local cluster.')
        cluster = LocalCluster()
        client = Client(cluster)
        logger.info('dask local cluster started.')

        read_gamma_image_delayed = delayed(read_gamma_image, pure=True)
        lazy_rslcs = [read_gamma_image_delayed(rslc,width,dtype='fcomplex') for rslc in rslcs]
        lazy_sim_orbs = [read_gamma_image_delayed(sim_orb,width,dtype='float') for sim_orb in sim_orbs]
        rslcs_data = [da.from_delayed(lazy_rslc,dtype= np.complex64,shape=(nlines,width)) for lazy_rslc in lazy_rslcs]
        rslcs_data = da.stack(rslcs_data,axis=2)
        sim_orb_data =[da.from_delayed(lazy_sim_orb,dtype=np.float32,shape=(nlines,width)) for lazy_sim_orb in lazy_sim_orbs]
        sim_orb_data = da.stack(sim_orb_data,axis=2)
        flatten_rslcs_data = da.exp(sim_orb_data*np.complex64(1j))*rslcs_data
        flatten_rslcs_data = flatten_rslcs_data.rechunk((az_chunk_size,width,1))
        logger.info(f'flattened rslcs shape: {flatten_rslcs_data.shape}')
        logger.info(f'flattened rslcs chunks: {flatten_rslcs_data.chunks}')
        _flatten_rslcs_data = flatten_rslcs_data.to_zarr(rslc_zarr,overwrite=True,compute=False)
        logger.info('computing graph setted. doing all the computing.')
        da.compute(_flatten_rslcs_data)
        logger.info('computing finished.')
        cluster.close()
        logger.info('dask cluster closed.')

# %% ../../nbs/CLI/load.ipynb 18
@call_parse
@log_args
def de_load_gamma_lat_lon(lt:str, # geocoding lookup table
                          rslc_par:str, # par file of the reference rslc
                          dem_par:str, # dem par
                          lat_zarr:str, # output, latitude zarr
                          lon_zarr:str, # output, longitude zarr
                          az_chunk_size:int=-1, # azimuth chunk size of lat and lon zarr, azimuth number of lines by default (one chunk)
                          log:str=None, # logfile, no log by default
):
    '''
    Function to load longitude and latitude from gamma binary format to zarr.
    Only DEM on EQA projection is supported now.
    '''
    logger = get_logger(logfile=log)
    geo_width = _geo_width_nlines(dem_par)[0]
    rdc_width, rdc_nlines = _rdc_width_nlines(rslc_par)
    logger.info(f'image shape: ({rdc_nlines},{rdc_width})')

    lat_data = zarr.open(lat_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)
    lon_data = zarr.open(lon_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        logger.info('run gamma command to generate longitude and latitude:')
        lt_invert = temp_dir/'rdc2geo'
        command = f'gc_map_inversion {lt} {geo_width} {str(lt_invert)} {rdc_width} {rdc_nlines} 2 2 16 16'
        logger.info('run command: ' + command)
        os.system(command)
        # pg.gc_map_inversion(lt,geo_width,lt_invert, rdc_width,rdc_nlines,
        #                     2,2,16,16)
        cor_lat,cor_lon,pos_lat,pos_lon = _cor_pos_dem(dem_par)
        rdc2geo = read_gamma_image(lt_invert,rdc_width,dtype='fcomplex')
        lat = cor_lat + pos_lat * rdc2geo.imag
        lon = cor_lon + pos_lon * rdc2geo.real
        logger.info('writing data.')
        lat_data[:] = lat
        lon_data[:] = lon
        logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 21
@call_parse
@log_args
def de_load_gamma_incidence_angle(inc:str, # incidence angle
                                  lt:str, # lookup table
                                  rslc_par:str, # par file of the reference rslc
                                  dem_par:str, # dem par
                                  inc_zarr:str, # output, latitude zarr
                                  az_chunk_size:int=-1, # azimuth chunk size of lat and lon zarr, azimuth number of lines by default (one chunk)
                                  log:str=None, # logfile, no log by default
):
    '''
    Load incidence angle from gamma binary format to zarr.
    The input incidence angle should be in EQA coordinate and generated by gc_map2 with inc_flg=1.
    '''
    logger = get_logger(logfile=log)
    geo_width = _geo_width_nlines(dem_par)[0]
    rdc_width, rdc_nlines = _rdc_width_nlines(rslc_par)
    logger.info(f'image shape: ({rdc_nlines},{rdc_width})')

    inc_data = zarr.open(inc_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        inc_rdc = temp_dir/'inc_rdc'
        logger.info('run gamma command to generate inc angle in range doppler coordinate:')
        command = f'geocode {lt} {inc} {geo_width} {str(inc_rdc)} {rdc_width} {rdc_nlines}'
        logger.info('run command: ' + command)
        os.system(command)
        logger.info('writing data.')
        inc_data[:] = read_gamma_image(inc_rdc,rdc_width,dtype='float')
        logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 24
@call_parse
@log_args
def de_load_gamma_hgt(hgt:str, # DEM in radar coordinate
                      rslc_par:str, # par file of the reference rslc
                      hgt_zarr:str, # output, latitude zarr
                      az_chunk_size:int=-1, # azimuth chunk size of lat and lon zarr, azimuth number of lines by default (one chunk)
                      log:str=None, # logfile, no log by default
):
    '''
    Load DEM in radar coordinate from gamma binary format to zarr.
    '''
    logger = get_logger(logfile=log)
    rdc_width, rdc_nlines = _rdc_width_nlines(rslc_par)
    logger.info(f'image shape: ({rdc_nlines},{rdc_width})')

    hgt_data = zarr.open(hgt_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        logger.info('writing data.')
        hgt_data[:] = read_gamma_image(hgt,rdc_width,dtype='float')
        logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 27
@call_parse
@log_args
def de_load_gamma_range(rslc_par:str, # par file of one rslc
                        range_zarr:str, # output, range distance zarr
                        az_chunk_size:int=-1, # azimuth chunk size of lat and lon zarr, azimuth number of lines by default (one chunk)
                        log:str=None, # logfile, no log by default
):
    '''
    Generate slant range distance and save to zarr.
    '''
    logger = get_logger(logfile=log)
    rdc_width, rdc_nlines = _rdc_width_nlines(rslc_par)
    logger.info(f'image shape: ({rdc_nlines},{rdc_width})')
    with open(rslc_par) as f:
        for line in f:
            if re.search('near_range_slc',line):
                rho0 = float(line.split()[1])
            if re.search('range_pixel_spacing',line):
                d_rho = float(line.split()[1])

    rho1d = np.arange(rdc_width)*d_rho+rho0
    rho2d = np.tile(rho1d,(rdc_nlines,1))
    range_data = zarr.open(range_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)
    logger.info('writing data.')
    range_data[:] = rho2d
    logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 30
@call_parse
@log_args
def de_load_gamma_metadata(rslc_dir:str, # # gamma rslc directory, the name of the rslc and their par files should be '????????.rslc' and '????????.rslc.par'
                           reference:str, # reference date, eg: '20200202'
                           meta_file:str, # text toml file for meta data
                           log:str=None, # logfile, no log by default
):
    '''
    Load necessary metadata into a toml file.
    '''
    meta = dict()
    rslc_dir = Path(rslc_dir)
    dates = []
    logger = get_logger(logfile=log)
    for par_file in sorted(rslc_dir.glob('*.rslc.par')):
        dates.append(par_file.name[:8])
    meta['dates'] = dates

    with open(rslc_dir/(reference+'.rslc.par')) as f:
        for line in f:
            if re.search('radar_frequency',line):
                radar_f = float(line.split()[1])
                rdr_wavelen = speed_of_light/radar_f
            if re.search('heading',line):
                heading = float(line.split()[1])
            if re.search('range_pixel_spacing',line):
                dr = float(line.split()[1])
            if re.search('azimuth_pixel_spacing',line):
                daz = float(line.split()[1])

    meta['radar_wavelength'] = rdr_wavelen
    meta['heading'] = heading
    meta['range_pixel_spacing'] = dr
    meta['azimuth_pixel_spacing'] = daz

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)
        slc_tab = temp_dir/'slc_tab'
        bperp = temp_dir/'bperp'
        itab = temp_dir/'itab'
        
        tab_content = ''
        for date in dates:
            tab_content += str(rslc_dir/(date+'.rslc'))
            tab_content += '      '
            tab_content += str(rslc_dir/(date+'.rslc.par'))
            tab_content += '\n'
        slc_tab.write_text(tab_content)

        command = f"base_calc {str(slc_tab)} {str(rslc_dir/(reference+'.rslc.par'))} {str(bperp)} {str(itab)} -"
        os.system(command)
        dat = pd.read_csv(bperp, sep='\s+', header=None)
        base=dat[3].to_numpy()
        base=base.astype(np.float32)

    meta['perpendicular_baseline'] = base

    with open(meta_file,'w') as f:
        a = toml.dump(meta,f,encoder=toml.TomlNumpyEncoder())
    logger.info('All meta data: \n'+a)
    logger.info('writing data in toml file.')
    logger.info('Done.')
