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

# %% auto 0
__all__ = ['read_gamma_image', 'write_gamma_image', 'read_gamma_plist', 'write_gamma_plist', 'de_load_gamma_flatten_rslc',
           'de_load_gamma_lat_lon_hgt', 'de_load_gamma_look_vector', 'de_load_gamma_range', 'de_load_gamma_metadata']

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

import numpy as np
from matplotlib import pyplot as plt
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'
    elif dtype == 'int':
        dt = '>i4'
    elif dtype =='double':
        dt = '>f8'
    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 10
def write_gamma_image(imag,path):
    imag = imag.astype(imag.dtype.newbyteorder('big'))
    imag.tofile(path)

# %% ../../nbs/CLI/load.ipynb 11
def read_gamma_plist(plist:str,dtype='int'):
    return read_gamma_image(plist,width=2,dtype=dtype)

# %% ../../nbs/CLI/load.ipynb 12
def write_gamma_plist(imag,path):
    return write_gamma_image(imag,path)

# %% ../../nbs/CLI/load.ipynb 14
@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)

    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 required data for flattened 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 > {str(temp_dir/"log")}'
            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 > {str(temp_dir/"log")}'
            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('gamma command finished.')
        logger.info('using dask to load data in gamma binary format to calculate flatten rslcs and save it 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.darr_info('flattened rslc', flatten_rslcs_data) 
        _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 20
@call_parse
@log_args
def de_load_gamma_lat_lon_hgt(diff_par:str, # geocoding diff_par,using the simulated image as reference
                              rslc_par:str, # par file of the reference rslc
                              dem_par:str, # dem par
                              hgt:str, # DEM in radar coordinate
                              lat_zarr:str, # output, latitude zarr
                              lon_zarr:str, # output, longitude zarr
                              hgt_zarr:str, # output, height 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.
    '''
    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.float64)
    lon_data = zarr.open(lon_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float64)
    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('run gamma command to generate longitude, latitude and height:')
        pt_i = np.arange(rdc_width,dtype=np.int32)
        pt_j = np.arange(rdc_nlines,dtype=np.int32)
        pt_ii,pt_jj = np.meshgrid(pt_i,pt_j)
        pt_ij = np.stack((pt_ii,pt_jj),axis=-1)
        write_gamma_plist(pt_ij,temp_dir/'plist')
        
        # no need to write phgt since it is identical to hgt if all point are selected.
        # hgt_data = read_gamma_image(hgt,width=rdc_width,dtype='float')
        # hgt_data = hgt_data.reshape(-1)
        # pg.write_point_data(hgt_data,temp_dir/'plist',temp_dir/'phgt')
        
        command = f"pt2geo {str(temp_dir/'plist')} - {rslc_par} - {hgt} {dem_par} {diff_par} 1 1 - - {str(temp_dir/'plat_lon')} {str(temp_dir/'phgt_wgs84')} > {str(temp_dir/'log')}"
        logger.info('run command: ' + command)
        os.system(command)
        logger.info('gamma command finished.')
        ptlonlat = read_gamma_plist(str(temp_dir/'plat_lon'),dtype='double')
        lon_data[:], lat_data[:] = ptlonlat[:,0].reshape(rdc_nlines,rdc_width),ptlonlat[:,1].reshape(rdc_nlines,rdc_width)
        hgt_data[:] = read_gamma_image(str(temp_dir/'phgt_wgs84'),width=rdc_width,dtype='float')
        logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 27
@call_parse
@log_args
def de_load_gamma_look_vector(theta:str, # elevation angle
                              phi:str, # orientation angle
                              lt:str, # lookup table
                              rslc_par:str, # par file of the reference rslc
                              dem_par:str, # dem par
                              theta_zarr:str, # output, elevation angle zarr
                              phi_zarr:str, # output, orientation angle 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 look vector (elevation angle and orientation angle) in map geometry
    from gamma binary format to look vector in radar geometry zarr file.
    The two input data should be generated with the `look_vector` gamma command.
    '''
    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})')

    theta_data = zarr.open(theta_zarr,mode='w',shape=(rdc_nlines,rdc_width),chunks = (az_chunk_size,rdc_width), dtype=np.float32)
    phi_data = zarr.open(phi_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)
        theta_rdc = temp_dir/'theta_rdc'
        logger.info('run gamma command to generate elevation angle in range doppler coordinate:')
        command = f'geocode {lt} {theta} {geo_width} {str(theta_rdc)} {rdc_width} {rdc_nlines} > {temp_dir/"log"}'
        logger.info('run command: ' + command)
        os.system(command)
        logger.info('gamma command finished.')
        logger.info('writing data.')
        theta_data[:] = read_gamma_image(theta_rdc,rdc_width,dtype='float')
        phi_rdc = temp_dir/'phi_rdc'
        logger.info('run gamma command to generate orientation angle in range doppler coordinate:')
        command = f'geocode {lt} {phi} {geo_width} {str(phi_rdc)} {rdc_width} {rdc_nlines} > {temp_dir/"log"}'
        logger.info('run command: ' + command)
        os.system(command)
        logger.info('gamma command finished.')
        logger.info('writing data.')
        phi_data[:] = read_gamma_image(phi_rdc,rdc_width,dtype='float')
        logger.info('Done.')

# %% ../../nbs/CLI/load.ipynb 33
@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])

    logger.info('Calculating slant range distance.')
    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 38
@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
                logger.info('Fetching randar wavelength')
            if re.search('heading',line):
                heading = float(line.split()[1])
                logger.info('Fetching heading angle')
            if re.search('range_pixel_spacing',line):
                dr = float(line.split()[1])
                logger.info('Fetching range pixel spacing')
            if re.search('azimuth_pixel_spacing',line):
                daz = float(line.split()[1])
                logger.info('fetching azimuth pixel spacing')

    meta['radar_wavelength'] = rdr_wavelen
    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)

        logger.info('Run gamma command to calculate baseline:')
        command = f"base_calc {str(slc_tab)} {str(rslc_dir/(reference+'.rslc.par'))} {str(bperp)} {str(itab)} - > {temp_dir/'log'}"
        logger.info('run command: ' + command)
        os.system(command)
        logger.info('gamma command finished.')
        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.')
