#!/usr/bin/env python3
# Alexey Pechnikov, Sep, 2021, https://github.com/mobigroup/gmtsar
from .SBAS_trans import SBAS_trans
from .PRM import PRM
from .tqdm_dask import tqdm_dask

class SBAS_topo_ra(SBAS_trans):

    def topo_ra(self, subswath=None, idec=2, jdec=2, n_jobs=-1, interactive=False):
        from scipy.spatial import cKDTree
        import dask
        import xarray as xr
        import numpy as np
        import os

        # trans.dat - file generated by llt_grid2rat (r a topo lon lat)"
        trans_dat = self.get_trans_dat(subswath)

        trans_blocks_extents = self.get_trans_dat_blocks_extents(subswath, n_jobs=n_jobs)

        # define topo_ra grid
        XMAX, yvalid, num_patch = self.PRM(subswath).get('num_rng_bins', 'num_valid_az', 'num_patches')
        YMAX = yvalid * num_patch
        #print ('DEBUG: XMAX', XMAX, 'YMAX', YMAX)
        # use center pixel GMT registration mode
        rngs = np.arange(1, XMAX+1, idec, dtype=np.int32)
        azis = np.arange(1, YMAX+1, jdec, dtype=np.int32)
        # do not use coordinate names y,x because the output grid saved as (y,y) in this case...
        azis = xr.DataArray(azis, dims=['a'], coords={'a': azis}).chunk(self.chunksize)
        rngs = xr.DataArray(rngs, dims=['r'], coords={'r': rngs}).chunk(self.chunksize)
        azi, rng = [da.chunk(self.chunksize) for da in xr.broadcast(azis, rngs)]

        def calc_topo_ra(azi, rng):
            # check thr arguments
            assert azi.shape == rng.shape, f'ERROR: {azi.shape} != {rng.shape}'

            # check the selected area bounds
            ymin, ymax, xmin, xmax = azi.min(), azi.max(), rng.min(), rng.max()
            #print ('ymin, ymax', ymin, ymax, 'xmin, xmax', xmin, xmax)
            # define corresponding trans_dat blocks
            ymask = (trans_blocks_extents[:,3]>=ymin-jdec)&(trans_blocks_extents[:,2]<=ymax+jdec)
            xmask = (trans_blocks_extents[:,5]>=xmin-idec)&(trans_blocks_extents[:,4]<=xmax+idec)
            blocks = trans_blocks_extents[ymask&xmask]
            #print ('trans_dat blocks', blocks.astype(int))

            blocks_azis = []
            blocks_rngs = []
            blocks_eles = []
            for iy, ix in blocks[:,:2].astype(int):
                #print ('iy, ix', iy, ix)
                block_azi = trans_dat.azi.data.blocks[iy, ix].persist()
                block_rng = trans_dat.rng.data.blocks[iy, ix].persist()
                block_ele = trans_dat.ele.data.blocks[iy, ix].persist()
                #print ('block_ele', block_ele.shape)

                blocks_azis.append(block_azi.reshape(-1))
                blocks_rngs.append(block_rng.reshape(-1))
                blocks_eles.append(block_ele.reshape(-1))

            # found the issue for 5 stitched scenes
            if len(blocks_azis) == 0:
                return np.nan*np.zeros(azi.shape)

            blocks_azis = np.concatenate(blocks_azis)
            blocks_rngs = np.concatenate(blocks_rngs)
            blocks_eles = np.concatenate(blocks_eles)

            # build index tree - dask arrays computes automatically
            source_yxs = np.stack([blocks_azis, blocks_rngs], axis=1)
            tree = cKDTree(source_yxs, compact_nodes=False, balanced_tree=False)

            # query the index tree
            target_yxs = np.stack([azi.reshape(-1), rng.reshape(-1)], axis=1)
            d, inds = tree.query(target_yxs, k = 1, workers=1)
            # compute dask array to prevent ineffective index loockup on it
            matrix = np.asarray(blocks_eles)[inds].reshape(azi.shape)
            return matrix

        # xarray wrapper
        topo = xr.apply_ufunc(
            calc_topo_ra,
            azi,
            rng,
            dask='parallelized',
            vectorize=False,
            output_dtypes=[np.float32],
        ).rename('topo_ra')

        if interactive:
            # do not flip vertically because it's returned as is without SBAS.get_topo_ra() function
            return topo

        # save to NetCDF file
        filename = self.get_filenames(subswath, None, 'topo_ra')
        if os.path.exists(filename):
            os.remove(filename)
        # flip vertically for GMTSAR compatibility reasons
        topo_ra = xr.DataArray(dask.array.flipud(topo), coords=topo.coords, name=topo.name)
        handler = topo_ra.to_netcdf(filename,
                                    encoding={'topo_ra': self.compression},
                                    engine=self.engine,
                                    compute=False)
        return handler

    def topo_ra_parallel(self, interactive=False):
        import dask

        # auto generate the trans.dat file
        self.trans_dat_parallel()

        # process all the subswaths
        subswaths = self.get_subswaths()
        delayeds = []
        for subswath in subswaths:
            delayed = self.topo_ra(subswath=subswath, interactive=interactive)
            delayeds.append(delayed)

        if not interactive:
            tqdm_dask(dask.persist(delayeds), desc='Radar Topography Computing')
        else:
            return delayeds[0] if len(delayeds)==1 else delayeds

    def get_topo_ra(self):
        #import numpy as np
        import xarray as xr
        import dask.array

        def func(topo):
            # flip vertically for GMTSAR compatibility reasons
            return xr.DataArray(dask.array.flipud(topo), coords=topo.coords, attrs=topo.attrs, name=topo.name)\
                   .rename({'a': 'y', 'r': 'x'})

        topos = self.open_grids(None, 'topo_ra', func=func)

        return topos[0] if len(topos)==1 else topos
